Add predict step accumulation
Author: sgugger
Pull request: yes

# What does this PR do?

Currently, the Trainer accumulates all predictions on the host (GPU or TPU) before gathering across all hosts (in case of distributed training) and moving back to the CPU. This can result in OOM errors when users have a big dataset (the model is already taking up a lot of space on the host) as highlighted in #7232. However moving the predictions to the CPU at each prediction step is also inefficient (particularly on TPU).

This PR aims at fixing the OOM problem while retaining efficiency by introducing a new training argument called `eval_accumulation_step`. If left untouched, the behavior is the same as right now (all predictions accumulated on the host and moved at the end of the prediction loop). If set to an int, the predictions are gathered and moved every `eval_accumulation_step`. This required some clever reorganization of the predictions (see the docstring of `DistributedTensorGatherer` for more details).

In passing I cleaned up the code related to gathering tensors across multiple hosts and fixed the issue of the `loss.item()` (big slow down to do that at every step on TPUs) and accumulated the losses the same way predictions and labels are. This still works for any number of outputs/labels of the model.

To check those changes did not break anything, I ran `test_trainer_distributed.py` on my local setup and created an equivalent for TPUs that I also ran (they both pass).

This slightly change Seq2SeqTrainer (since we don't want the `loss.item()`) so cc @patil-suraj I don't think this should break anything in it.

<!-- Remove if not applicable -->

Fixes #7232


Comments:
Thanks so much @sgugger - you're a legend! 🙇 

> # What does this PR do?
> Currently, the Trainer accumulates all predictions on the host (GPU or TPU) before gathering across all hosts (in case of distributed training) and moving back to the CPU. This can result in OOM errors when users have a big dataset (the model is already taking up a lot of space on the host) as highlighted in #7232. However moving the predictions to the CPU at each prediction step is also inefficient (particularly on TPU).
> 
> This PR aims at fixing the OOM problem while retaining efficiency by introducing a new training argument called `eval_accumulation_step`. If left untouched, the behavior is the same as right now (all predictions accumulated on the host and moved at the end of the prediction loop). If set to an int, the predictions are gathered and moved every `eval_accumulation_step`. This required some clever reorganization of the predictions (see the docstring of `DistributedTensorGatherer` for more details).
> 
> In passing I cleaned up the code related to gathering tensors across multiple hosts and fixed the issue of the `loss.item()` (big slow down to do that at every step on TPUs) and accumulated the losses the same way predictions and labels are. This still works for any number of outputs/labels of the model.
> 
> To check those changes did not break anything, I ran `test_trainer_distributed.py` on my local setup and created an equivalent for TPUs that I also ran (they both pass).
> 
> This slightly change Seq2SeqTrainer (since we don't want the `loss.item()`) so cc @patil-suraj I don't think this should break anything in it.
> 
> Fixes #7232

Thanks so much @sgugger 
`eval_accumulation_steps` for the argument name and not `eval_accumulation_step` 😉