Preprocess/transform logits before caching them for computing metrics.
# 🚀 Feature request

I think it'd be nice to have a simple way to preprocess the logits before caching them for computing metrics.

## Motivation

When the `Trainer` `compute_metrics` are set, during evaluation the logits are accumulated (some in GPU memory, for `args.eval_accumulation_steps` steps; all in RAM). For some models, it will almost certainly lead to out of memory problems.

For instance, for a language model, this means storing in RAM a tensor of size [eval ds size, sequence length, vocab size].

In many cases, what is needed to compute metrics is just some reduction of the logits. For example: `logits.argmax(dim=-1)`.

I know I can subclass `Trainer` for this and redefine `evaluation_loop`, just wanted to know if you'd consider a more generic solution that prevents everyone that needs the feature from duplicating the rest of the code of `evaluation_loop`. I've seen more people running into the same issue. For instance:

https://github.com/huggingface/transformers/issues/8476
https://discuss.huggingface.co/t/cuda-out-of-memory-when-using-trainer-with-compute-metrics/2941
https://discuss.huggingface.co/t/cuda-out-of-memory-during-evaluation-but-training-is-fine/1783/4

## Your contribution

I was thinking about something like adding a `preprocess_logits_for_metrics` parameter to `TrainingArguments` of type Callable

If you don't set the parameter, the default is None and everything would work as always. If you set it, the logits are passed to `args.preprocess_logits_for_metrics` and its output is what's cached.

The main modification would be this in `Trainer.evaluation_loop`:

```
# Update containers on host
...
 if logits is not None:
    logits = self._pad_across_processes(logits)
    logits = self._nested_gather(logits)
    if self.args.preprocess_logits_for_metrics is not None:
        logits = self.args.preprocess_logits_for_metrics(logits)
    preds_host = logits if preds_host is None else nested_concat(preds_host, logits, padding_index=-100)
```

Do you think it's worth it? If you do, I can submit a PR.

I tag @sgugger because I think he's worked quite a lot with the training loop, but I'm open to receive feedback from anyone.


I think it would be a valuable addition, as you describe the problematic situation very well, when someone wants to compute perplexity with a language model having a very large vocab size, for instance.

The `TrainingArguments` can't have a new argument of type callable, but I think we could have a new argument in the init `preprocess_logits_for_metrics`.

I'm happy to review a PR for this, and if you could show inside how to use it in the examples `run_clm` or `run_mlm` to get the perplexity at each evaluation without getting OOM, that would be a very compelling argument for this new API!

cc @LysandreJik for info.