evaluating with trainer.py with TPU results in sudden RAM spike and crash
Author: LeonieWeissweiler
Pull request: no
Labels: name: wontfix

# 🐛 Bug

## Information

Model I am using (Bert, XLNet ...): Roberta with the official how to train from scratch example

Language I am using the model on (English, Chinese ...): Esperanto

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [x] my own modified scripts: (give details below)

## To reproduce

Steps to reproduce the behavior:

1. follow the official how to train with esperanto example
2. define a second dataset for testing, both datasets are TextDatasets, the train set is around 300MB, the test set around 3MB
3. run trainer.train() with evaluate_during_training

The evaluation loop runs a few examples before crashing the colab session due to an unknown reason. I have managed to increase the number of examples from 2 to 6 by reducing the per_device_eval_batch_size. My eval dataset only contains 31 examples. When checking the colab logs after crashing, it seems that trainer.py tried to allocate 4GB of memory before crashing, which seems unrealistic to me given the size of the datasets.

Specifically (I think this is the relevant part, please correct me if I'm wrong):

2020-06-14 21:45:53.187694: E tensorflow/compiler/xla/xla_client/xla_util.cc:76] (1) Resource exhausted: Attempting to reserve 2.78G at the bottom of memory. That was not possible. There are 2.75G free, 0B reserved, and 2.75G reservable.

When switching to GPU, I encounter a similar issue, but this time with a complete stack trace because this doesn't crash the runtime:

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-14-10e96b9a36a3> in <module>()
----> 1 trainer.train()
      2 trainer.save_model("./drive/My Drive/models/roberta/output")

2 frames
<ipython-input-11-34c6e8e4c40d> in train(self, model_path)
    507 
    508                         if self.args.evaluate_during_training:
--> 509                             self.evaluate()
    510 
    511                     if self.args.save_steps > 0 and self.global_step % self.args.save_steps == 0:

<ipython-input-11-34c6e8e4c40d> in evaluate(self, eval_dataset, prediction_loss_only)
    706         eval_dataloader = self.get_eval_dataloader(eval_dataset)
    707         print(2)
--> 708         output = self._prediction_loop(eval_dataloader, description="Evaluation")
    709         print(3)
    710         self._log(output.metrics)

<ipython-input-11-34c6e8e4c40d> in _prediction_loop(self, dataloader, description, prediction_loss_only)
    775                     preds = logits.detach()
    776                 else:
--> 777                     preds = torch.cat((preds, logits.detach()), dim=0)
    778                 if inputs.get("labels") is not None:
    779                     if label_ids is None:

RuntimeError: CUDA out of memory. Tried to allocate 6.75 GiB (GPU 0; 15.90 GiB total capacity; 7.99 GiB already allocated; 6.40 GiB free; 8.81 GiB reserved in total by PyTorch)

## Expected behavior

I expected the eval loss or error to be computed. 

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->
     
- `transformers` version: 2.11.0
- Platform: colab
- Python version: 3.6
- Using GPU in script?: no, TPU with xla
- Using distributed or parallel set-up in script?: in theory, but only 1 TPU core available


Comments:
This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.


Any update on this? I've got a similar issue