[trainer] fix nan in full-fp16 label_smoothing eval
Author: stas00
Pull request: yes

This PR fixes the issue of getting NaN eval loss with any inference that uses full fp16 model - which is the case with deepspeed or when `--fp16_full_eval` is passed.

The problem is that `log_probs.sum` runs over 30-50K of numbers overflows easily in fp16, so this PR switches it to fp32 internally. Which surprisingly requires almost no extra memory. As the conversion happens on the hardware level and we only need an extra `2 bytes * batch_size` of additional memory.

Here is some data showing the the metrics remain the same after this fix:

```
export BS=16; rm -r output_dir; PYTHONPATH=src USE_TF=0 CUDA_VISIBLE_DEVICES=0 \
./examples/seq2seq/run_translation.py --model_name_or_path t5-small --output_dir /tmp/zero3 \
--overwrite_output_dir --max_train_samples 10 --max_val_samples 100 --max_source_length 12 \
--max_target_length 128 --val_max_target_length 128 --do_train --num_train_epochs 1 \
--per_device_train_batch_size 2 --learning_rate 3e-3 --warmup_steps 8 --predict_with_generate \
--logging_steps 0 --save_steps 2 --eval_steps 1 --group_by_length --adafactor --dataset_name wmt16 \
--dataset_config ro-en --source_lang en --target_lang ro \
--source_prefix "translate English to Romanian: " --do_eval --label_smoothing 0.1

***** train metrics *****
  epoch                      =    1.0
  init_mem_cpu_alloc_delta   =    2MB
  init_mem_cpu_peaked_delta  =    0MB
  init_mem_gpu_alloc_delta   =  230MB
  init_mem_gpu_peaked_delta  =    0MB
  train_mem_cpu_alloc_delta  =   60MB
  train_mem_cpu_peaked_delta =   63MB
  train_mem_gpu_alloc_delta  =  231MB
  train_mem_gpu_peaked_delta =  194MB
  train_runtime              = 7.7162
  train_samples              =     10
  train_samples_per_second   =  0.648
  
***** eval metrics *****
  epoch                     =    1.0
  eval_bleu                 = 2.4612
  eval_gen_len              =  18.53
  eval_loss                 =  5.017
  eval_mem_cpu_alloc_delta  =    0MB
  eval_mem_cpu_peaked_delta =    0MB
  eval_mem_gpu_alloc_delta  =    0MB
  eval_mem_gpu_peaked_delta =  244MB
  eval_runtime              = 4.6481
  eval_samples              =    100
  eval_samples_per_second   = 21.514
```

now let's add `--fp16_full_eval`, which before this PR leads to ` eval_loss = nan`
```
export BS=16; rm -r output_dir; PYTHONPATH=src USE_TF=0 CUDA_VISIBLE_DEVICES=0 \
./examples/seq2seq/run_translation.py --model_name_or_path t5-small --output_dir /tmp/zero3 \
--overwrite_output_dir --max_train_samples 10 --max_val_samples 100 --max_source_length 12 \
--max_target_length 128 --val_max_target_length 128 --do_train --num_train_epochs 1 \
--per_device_train_batch_size 2 --learning_rate 3e-3 --warmup_steps 8 --predict_with_generate \
--logging_steps 0 --save_steps 2 --eval_steps 1 --group_by_length --adafactor --dataset_name wmt16 \
--dataset_config ro-en --source_lang en --target_lang ro \
--source_prefix "translate English to Romanian: " --do_eval --label_smoothing 0.1 \
--fp16_full_eval

***** train metrics *****
  epoch                      =    1.0
  init_mem_cpu_alloc_delta   =    2MB
  init_mem_cpu_peaked_delta  =    0MB
  init_mem_gpu_alloc_delta   =  230MB
  init_mem_gpu_peaked_delta  =    0MB
  train_mem_cpu_alloc_delta  =   60MB
  train_mem_cpu_peaked_delta =   63MB
  train_mem_gpu_alloc_delta  =  231MB
  train_mem_gpu_peaked_delta =  194MB
  train_runtime              = 7.1477
  train_samples              =     10
  train_samples_per_second   =    0.7

***** eval metrics *****
  epoch                     =    1.0
  eval_bleu                 = 2.4612
  eval_gen_len              =  18.53
  eval_loss                 = 5.0168
  eval_mem_cpu_alloc_delta  =    0MB
  eval_mem_cpu_peaked_delta =    0MB
  eval_mem_gpu_alloc_delta  = -231MB
  eval_mem_gpu_peaked_delta =  262MB
  eval_runtime              = 6.0125
  eval_samples              =    100
  eval_samples_per_second   = 16.632
```
`eval_loss` is off by 0.0002.

I spent quite some time trying to find where to add a test, but it's a tricky situation where the input has to be pretty huge. I remember seeing it in some deepspeed tests, but I can't find at the moment, currently all tests return a normal number.

One interesting things I noticed is that `--fp16_full_eval` makes eval slower by 20-25% which is strange, but I tested this PR has no impact on the speed. I file a separate issue about it https://github.com/huggingface/transformers/issues/10816

Fixes: https://github.com/huggingface/transformers/issues/10674

@sgugger, @LysandreJik 

Comments:
Hi @stas00 
I tested this PR and for me this becomes really slower with mt5-small model after adding the modifications in this PR, here is the command I run, I am using transformer=4.4.2. I will be grateful to your expert knowledge to have the speed issue also fixed. Thank you very much for the incredible jobs you do. 

`deepspeed run_translation.py     --model_name_or_path google/mt5-small     --do_train     --do_eval     --source_lang en     --target_lang ro     --dataset_name wmt16     --dataset_config_name ro-en     --output_dir test/tst-t1ranslatieeeon     --per_device_train_batch_size=4     --per_device_eval_batch_size=4          --predict_with_generate --max_train_samples 100 --fp16 --deepspeed ds_config.json --max_val_samples 100 --logging_step 10
`

I can't possibly see how this PR could impact the speed since it changes the label_smoother and your command line doesn't have `--label_smoothing 0.1` so the modified code in this PR doesn't get to run.

That said when you use this PR you're in a way using `master`, so perhaps you were testing with some other `transformers` version before and noticed a regression in `master`. Try to test with whatever version you were using before and then retest with the `master` branch and see whether you can reproduce your issue. 

If it is, do you know how to use `git bisect`? You can then in a matter of a few runs find out the commit that impacted the performance. If you can't figure it out just give me the last good transformers version and I will help you from there.

-----

Also you're not telling me what's inside `ds_config.json` - I assume it's zero2 configuration. zero3 isn't quite ready yet.