Store Predictions on CPU in Every Prediction Iteration (Trainer)
# 🚀 Feature request

Store Predictions on CPU in Every Prediction Iteration (Trainer)

## Motivation

Currently, in  [`Trainer._prediction_loop`,](https://github.com/huggingface/transformers/blob/33d7506ea10ca92886fd1bb3b5306a1a720c58fe/src/transformers/trainer.py#L785) the predictions (logits) of the model is stored on GPU/TPU in each iteration. After all iterations are finished, they will be concatenated together and sent to CPU. In this way, the GPU/TPU memory usage will increase linearly during prediction. If the test set is very large, there may be insufficient GPU/TPU memory to finish the prediction phase. To save GPU/TPU memory and allow larger-scale inference, the trainer should instead send a batch of prediction to CPU after each iteration.

However, if we do inference on multiple devices, we would need a `distributed_concat` function to aggregate all predictions from all devices. If the predictions are already stored on CPU, we can no longer aggregate them using NCCL. Therefore, this issue is unsolvable unless we introduce another CPU-based distributed communication mechanism. Still, we can at least solve it in a single-GPU scenario.

## Your contribution

If approved, I can write codes to solve this problem when the inference is not running in the distributed mode.
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.
