trainer.evaluate() aggregates predictions on GPU and causes CUDA out of memory issues for large datasets
## Environment info
- `transformers` version: 3.1.0
- Platform: Linux-4.15.0-112-generic-x86_64-with-glibc2.10
- Python version: 3.8.5
- PyTorch version (GPU?): 1.6.0 (True)
- Tensorflow version (GPU?): not installed (NA)
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: No

### Who can help
 
Trainer: @sgugger

## Information

Model I am using (Bert, XLNet ...): Bert

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [X] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [X] my own task or dataset: (give details below)

## To reproduce

Steps to reproduce the behavior:

1. Use the Trainer for evaluation (`.evaluate()`, `.predict()`) on the GPU with BERT with a large evaluation DataSet where the size of the returned prediction Tensors + Model exceed GPU RAM. (In my case I had an evaluation dataset of 469,530 sentences).
2. Trainer will crash with a CUDA Memory Exception

## Expected behavior

- I would expect the predictions in `predict()` or `evaluate()` from each step would be moved to the CPU device (off the GPU) and then concatenated later. However, the tensors are concatenated while still be on the GPU device, and only converted to CPU numpy arrays after the whole dataset has been predicted/evaluated. 

This means that for large evaluation datasets you'll run out of CUDA memory.

It also makes it difficult to pick the batch size to optimize the batch size for the GPU, as you need to allow space for not only the model and inputs, but also all the predictions, which can add up when dealing with large evaluation datasets.

- ie. The problem in the trainer code is that the predictions stay on the GPU [here](https://github.com/huggingface/transformers/blob/67d9fc50d917c63cf67281106214e1d9ae018dff/src/transformers/trainer.py#L1315)
- These tensors get concatenated but stay on the GPU [here](https://github.com/huggingface/transformers/blob/67d9fc50d917c63cf67281106214e1d9ae018dff/src/transformers/trainer_utils.py#L143)
- and then the predictions eventually end up on the CPU [here](https://github.com/huggingface/transformers/blob/67d9fc50d917c63cf67281106214e1d9ae018dff/src/transformers/trainer.py#L1343-L1347)

What this means, is that for larger evaluation datasets, *all* the predictions stay on the GPU, which is a function of how long your evaluation dataset is, and that you often run out of GPU RAM.

A work around is to something like this and run the loop yourself, and predict by batch:

``` py
    preds = []
    for i in tqdm(range(0, len(ds_valid), step)):
        ds_valid = KeywordDataset(tokenizer, df_valid[i:i+step], targets)
        batch_preds = trainer.predict(ds_valid)
        preds.append(batch_preds)
        batch_accuracy = (batch_preds.predictions.argmax(-1) == batch_preds.label_ids).mean()

    np_preds = np.concatenate([pred.predictions for pred in preds], axis=0)
    np_label_ids = np.concatenate([pred.label_ids for pred in preds], axis=0)
    acc = (np_preds.argmax(-1) == np_label_ids).mean()
    print('eval_accuracy = ', acc)
``` 

It would be nice if it defaulted the cpu (so you don't have to worry about it), and that there is another trainer argument to give the device to aggregate predictions that can be used to override that behaviour if it's important for certain use cases.

Thanks for the amazing work on an amazing library. Working with transformers has never been easier due to the hard work of your team!
Yes that is a known issue. The workaround you suggest does not work for distributed training, as the tensors need to be kept on GPU/TPU to be reduced together. Our plans to address this problem are to use the same kind of code as the metrics in `datasets`, but this is a bit of work and will take us a few weeks to fix.

In the meantime, I suggest you use the `predict` method and feed your dataset in slices there to avoid the OOM (then concatenate them on the CPU before running your metrics) or if you don't care about distributed eval, you can just subclass `Trainer` and replace the problematic code you identified by your own.<|||||>@eugeneware for context, what kind of task are you training on, and what's the approximate size/dimensions of the aggregated predictions tensor?<|||||>@sgugger thanks for the update and getting back to me. I'm currently running just a single GPU for this task, so the subclassing will work OK at the moment. Given your work in the fast.ai trainer, I'm excited to see where the `Learner` goes! :-)

---

@julien-c - I'm training a classifier with a large number of output classes (9997 classes).

I have 469530 items in my validation dataset.

So, by math, it would need to accumulate 469530 * 9997 * 4 = 17.5 GB of VRAM at fp32 or 8.8GB at fp16?

I'm training on a Titan RTX with 24GB of RAM, and I've found with the `bert-base-uncased` model that evaluation seems to die a bit over 1/2 way through training.

---

For others who might stumble across this issue, and are not using distributed training,  here's my subclass of the `Trainer` class - I just changed the last line to return the tensors moved onto the `eval_device` argument which I'm defaulting to the cpu:

``` py
class CustomTrainer(Trainer):
    def __init__(self, *args, eval_device=torch.device('cpu'), **kwargs):
        super().__init__(*args, **kwargs)
        self.eval_device = eval_device

    def prediction_step(
        self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]], prediction_loss_only: bool
    ) -> Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]:
        has_labels = any(inputs.get(k) is not None for k in ["labels", "lm_labels", "masked_lm_labels"])

        inputs = self._prepare_inputs(inputs)

        with torch.no_grad():
            outputs = model(**inputs)
            if has_labels:
                loss, logits = outputs[:2]
                loss = loss.mean().item()
            else:
                loss = None
                logits = outputs[0]
            if self.args.past_index >= 0:
                self._past = outputs[self.args.past_index if has_labels else self.args.past_index - 1]

        if prediction_loss_only:
            return (loss, None, None)

        labels = inputs.get("labels")
        if labels is not None:
            labels = labels.detach()
        # move tensors to evaluation device
        ret = (loss, logits.detach().to(self.eval_device), labels.to(self.eval_device))
        return ret
```<|||||>That *is* a large number of output classes.

I hope you post it to huggingface.co after training so that it breaks my inference widget 😅
https://huggingface.co/roberta-large-mnli?text=The+dog+was+lost.+Nobody+lost+any+animal<|||||>Thanks for creating this issue @eugeneware
I just hit this issue too, but in a distributed setting.

As a quick fix, I thought about calling `distributed_concat` after every batch and keep `preds`/`label_ids` in cpu. My env has much larger CPU RAM compared to GPU RAM. 

@sgugger Could you point to the corresponding part in datasets? Curious about this approach too.
> Our plans to address this problem are to use the same kind of code as the metrics in datasets, but this is a bit of work and will take us a few weeks to fix.



<|||||>There is a workaround (chunking your dataset in smaller parts) which is why we didn't implement a quick fix.
I'll start working on the right fix next week and link the PR here once it's ready @usuyama, so you can see the corresponding part in datasets (I have not yet investigated the exact part in datasets so can't answer your question right now, I've just been told it should be possible ;-) )<|||||>In case this solution is not suitable for someone like me, I had the same problem: `out of memory` error using cuda. But for my dataset, I could stay on the CPU for training. So I wrote my own training function and converted the dataset to a pytorch dataloader and it's done. Then I can totally control the processor 😃.<|||||>Is this problem solved?<|||||>I am still facing the error, Is there some code which would not get this error if executed?