[trainer] loss = NaN with label_smoothing and full-fp16 eval
Author: stas00
Pull request: no
Labels: name: Good First Issue

name: Good Second Issue

It looks like our `--label_smoothing_factor` Trainer's feature doesn't handle fp16 well. It's a problem with the deepspeed zero3 I'm integrating right now, since it evals in fp16, but also can be reproduced with the recently added `--fp16_full_eval` trainer option.

To reproduce:
```
export BS=16; rm -r output_dir; PYTHONPATH=src USE_TF=0 CUDA_VISIBLE_DEVICES=0 python examples/seq2seq/run_seq2seq.py --model_name_or_path t5-small --output_dir output_dir --adam_eps 1e-06 --do_eval --evaluation_strategy=steps  --label_smoothing 0.1 --learning_rate 3e-5 --logging_first_step --logging_steps 1000 --max_source_length 128 --max_target_length 128 --num_train_epochs 1 --overwrite_output_dir --per_device_eval_batch_size $BS --predict_with_generate --eval_steps 25000  --sortish_sampler --task translation_en_to_ro  --val_max_target_length 128 --warmup_steps 500  --max_val_samples 500 --dataset_name wmt16 --dataset_config "ro-en" --source_prefix "translate English to Romanian: " --fp16_full_eval
```
```
***** eval metrics *****
  eval_bleu                 = 24.1257
  eval_gen_len              =  39.554
  eval_loss                 =     nan
  eval_mem_cpu_alloc_delta  =    56MB
  eval_mem_cpu_peaked_delta =     0MB
  eval_mem_gpu_alloc_delta  =   116MB
  eval_mem_gpu_peaked_delta =   374MB
  eval_runtime              = 25.3246
  eval_samples              =     500
  eval_samples_per_second   =  19.744
  init_mem_cpu_alloc_delta  =     2MB
  init_mem_cpu_peaked_delta =     0MB
  init_mem_gpu_alloc_delta  =     0MB
  init_mem_gpu_peaked_delta =     0MB
```

If someone in the community would like to have a look at solving this puzzle, please refer to the discussion of this Issue.

Basically, we would like to try to find a way to perform label smoothing under full fp16 while finding a way to handle NaNs so that the final loss is not a NaN. 

And for the reference value running the same script w/o `--fp16_full_eval` should give you the "golden" `eval_loss` - i.e. ideally it should be about the same with  `--fp16_full_eval` (if possible that is).

Thank you!

@sgugger 

Comments:
I am unsure what you want me to fix. Yes evaluation in FP16 is not as precise as evaluation in FP32 and results in NaNs more easily, in particular in the loss, that's precisely the reason I was reluctant to add the --fp16_full_eval option.

well, first of all I'm reporting this behavior.

Secondly, it looks like DeepSpeed Zero3 for now requires that we continue using the fp16 model for inference. We are discussing other possibilities but they aren't there yet. This is because of all the hooks that get installed into the model during training. But if we were to remove all those hooks suddenly the model won't fit into the gpu memory - if it was spread out over multiple gpus in the first place. 

Which is why I added ` --fp16_full_eval` in first place. To enable fitting the model onto the gpu memory during inference if it fit during training (with fp32 master weights being offloaded to cpu). e.g. fitting t5-11b (45GB at fp32) into 40GB gpu. Can train with DeepSpeed zero2, so should be able to do inference too.

Finally, according to DeepSpeed devs a model trained in mixed precision shouldn't need to be put back into fp32 and should provide similar results only slightly imprecise. 

we will need a way to get the fp32 model out of DeepSpeed training https://github.com/microsoft/DeepSpeed/issues/800

I understand all of that, and maybe a fine-tuned model in mixed precision will have better results that don't get any Nan losses on the evaluation dataset, but it only takes one sample of the whole evaluation dataset with a slightly bigger loss than usual to get one Nan that will then drive the whole evaluation loss to NaN. And with label smoothing enabled it only takes one log_probs over all the possibilities to get to that.

Not sure it really matters since you still have proper metrics (the predictions are not driven to Nan, just the loss). Maybe we could try a flag to deactivate label smoothing when running evaluation, but then the loss wouldn't be comparable to the training loss, so not sure if it would really be useful either.


Thank you for elucidating the situation and suggesting workarounds, @sgugger 

Would it be of a useful value if we approximated the nans with something that would lead to a non-nan loss? since NaN can come from a variety of combinations of inf/0 operations, do we know which one is it? and then perhaps pick a corresponding substitution that might lead to a sufficiently good estimate?



Yes, we could maybe track when it arises and where it comes from exactly as a first step. Won't have time to dive into this in the near future, but if someone wants to tackle that issue to that effect and report here, that would be awesome!

I'm interested in taking a stab on this!

Awesome! Thank you, @vladdy

Just be aware that the following PR should be merged shortly: https://github.com/huggingface/transformers/pull/10611
and so the script in the reproduction line will most likely be renamed to `run_translation.py` and will slightly change the cl args.


@vladdy, while you research this - it'd be great to understand the cause of NaNs - so if you discover which operation leads to it please do share. Thank you!

Before falling asleep this idea came to me, tried it this morning and it worked:

```
--- a/src/transformers/trainer_pt_utils.py
+++ b/src/transformers/trainer_pt_utils.py
@@ -390,7 +390,9 @@ class LabelSmoother:

     def __call__(self, model_output, labels):
         logits = model_output["logits"] if isinstance(model_output, dict) else model_output[0]
+        #logits = logits.to(dtype=torch.float32)
         log_probs = -torch.nn.functional.log_softmax(logits, dim=-1)
+        log_probs = log_probs.to(dtype=torch.float32)
         if labels.dim() == log_probs.dim() - 1:
             labels = labels.unsqueeze(-1)
```

Basically, flip `logits` or `log_probs` back to fp32 and the problem goes away.

So the problem here has nothing to do with full fp16 inference, but with how label smoothing is calculated.

The issue comes from:
```
        smoothed_loss = log_probs.sum(dim=-1, keepdim=True)
```
in fp32, the return values are huge:
```
         [ 863637.3750],
         [ 864242.0000],
         [ 865449.0000],
         [ 866092.9375],
         [ 867702.4375],
```
and in fp16 these turn `inf`.

So either:
1. we do what I proposed on top of this comment, which will double the size of the `log_probs` tensor (4 times if we apply it to `logits`, rather than `log_probs`) . This will of course depend on the size of the dictionary and `target_max_len` - so say:
```
bs * max_len * dict_size * 2 more bytes
32 * 128 * 32000 * 2 = 262MB - huge!
```
2. or we change the calculation to scale down huge numbers back to a numerical range where `sum` over fp16 numbers doesn't overflow. Note that `smoothed_loss` does another `sum` towards the end which would definitely blow things up.
3. same as (1) but switch the calculations to `.cpu()` - a bit slower but no extra gpu memory will be required.

Of course number 2 is a better solution since it doesn't require much more memory to solve this problem and will require a change in algorithm to avoid going into huge numbers.

The other question is: do we deal with the label smoother separately or do we have other parts which may be affected in which case we should change the logits back to fp32 when prediction has completed. But as explained above this will come at a large gpu memory cost.

Yeah, I came to a similar conclusion regarding the cause and when I wanted to post an update on it,  I saw @stas00's response above.

If I'm not wrong, Apex tried to solve [this compatibility issue between fp16 and losses](https://github.com/NVIDIA/apex/tree/a109f856840ebb3ff5578e0bddfc4cffd4b96ed0/apex/fp16_utils), but I'm not sure how much of that could  be reused in addition to already stated options. @stas00, please let me know if you want to continue driving this and I'll try to find some other issue for my contribution.

@vladdy, by all means please continue, I was just sharing what I have discovered and calculated that this won't be an efficient solution memory requirement-wise. And I now have a better understanding of where NaN came from.

As you're suggesting the most efficient generic solution would be around loss scaling. We are doing it already during the training, so this is just some of the same for label smoothing.

But we definitely don't want to depend on apex for this. 

I haven't looked closely but I think the idea is to scale the `log_probs` into a much smaller size, while ensuring that the scaled numbers and the sum of 30-60k elements remain within the dynamic range of fp16 (plus there is one more sum of sums at the end!). If we were to implement it in a non-vectorized way it'd be the simplest to create an fp32 variable and add the fp16 bits to it, so it won't overflow. It won't take any extra memory, but that won't be efficient speed wise.

And of course perhaps you can think of other solutions. Anything that doesn't require an extra GPU memory to perform label smoothing is goodness.

Perhaps pytorch has some ready-made solutions too...

Any progress on this, @vladdy? I have one idea that may work, but would love to hear what you have discovered in your research.

@stas00, I have not found anything better than switching to fp32 for that operation. The rest of the approaches appear to be more complicated or not as generic as I think we want them to be. What idea did you have in mind?

Thank you for sharing the results of your research, @vladdy. So the need is the same, it's all about how to do it efficiently and not defeat the purpose of keeping things at fp16.

My discovery was switching to fp32 only for the aggregate and have it done by pytorch on the hardware level:
```
diff --git a/src/transformers/trainer_pt_utils.py b/src/transformers/trainer_pt_utils.py
index ae8e24949..c2071733c 100644
--- a/src/transformers/trainer_pt_utils.py
+++ b/src/transformers/trainer_pt_utils.py
@@ -399,7 +399,8 @@ class LabelSmoother:
         # will ignore them in any case.
         labels.clamp_min_(0)
         nll_loss = log_probs.gather(dim=-1, index=labels)
-        smoothed_loss = log_probs.sum(dim=-1, keepdim=True)
+        # works for fp16 input tensor too, by internally upcasting it to fp32
+        smoothed_loss = log_probs.sum(dim=-1, keepdim=True, dtype=torch.float32)

         nll_loss.masked_fill_(padding_mask, 0.0)
         smoothed_loss.masked_fill_(padding_mask, 0.0)
```
I was shocked that it took ~0 extra memory over pure fp16 solution (that is not even extra peak memory!) - that is it was doing the conversion to fp32 on the hardware level. At least that was so on my hardware - it might be not so on an older gpu that doesn't support fp16 natively. 

This is pretty amazing that it can do that!

I came to this idea while researching bfloat16 where its aggregate operations too require fp32 aggregates - so I thought why not try the same for our case and it seems to work. Similarly sometimes they use fp64 for fp32 inputs if they are too big, but I don't think we run into those things here.

What do you think? 


I think, this simple solution makes sense to be applied as it is also generic enough to cover all the cases. I doubt it is possible to find a better approach within short time and it does not appear it is necessary to spend more time on this (at least, for now). Feel free to do the PR as you offered it!

Dear @stas00 
Thank you very much for filing this issue. I am training mt5-small model and with deepspeed without label smoothing, I am getting NaNs, so far could not managed to fix it. I greatly appreciate your suggestions on this. if you think this can be appropriate I will open up a separate issue for mt5 model getting NaNs with deepspeed, and if not I follow this issue. Thank you very much

@vladdy, thank you for doing your research and validating my suggestion.

@dorost1234, the PR is here: https://github.com/huggingface/transformers/pull/10815 or you can just change it manually for your test - it's just one line.
Please let me know if it fixes your problem, if it's about eval_loss. If it doesn't, or it's about something else - then yes please open a separate issue. Based on your comments elsewhere the issue about mt5 and NaNs already exists, but not with deepspeed so definitely open one. Perhaps the Deepspeed team has some insights about this situation.



Thank you so much @stas00 for the great response, I applied your PR and with deepspeed now mt5-small for me is not getting nan anymore, this is an incredible job you are doing, thanks a lot, I still getting nans with mt5-small with fp16, even after your PR, for this I made a separate issue here https://github.com/huggingface/transformers/issues/10819 I did not tag you since with deepspeed with your applied magic PR it is not getting nans so far, while still If you have time to give me an advice I would be really grateful. 

@stas00 I tested mt5-small with run_translation.py model without this PR this also works fine without nans, if one does not use smoothing, with this PR this becomes much slower for me with deepspeed. is there a way to keep the speed as great as before?

Please post the exact command lines that you're referring to.

As I wrote in https://github.com/huggingface/transformers/pull/10815 I definitely can see a 25% slowdown when enabling --fp16_full_eval and opened an issue about it https://github.com/huggingface/transformers/issues/10816 

I don't see any speed difference with https://github.com/huggingface/transformers/pull/10815 w/o deepspeed, so once you show me what command line use then I can test.

edit: Oh, I see you posted them in https://github.com/huggingface/transformers/issues/10819 - all is good then - I can test now.