[trainer] loss = NaN with label_smoothing and full-fp16 eval

[stas00]

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

[sgugger]

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.

[stas00]

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 deepspeedai/DeepSpeed#800

[sgugger]

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.

[stas00]

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?

[sgugger]

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!

[vladdy]

I'm interested in taking a stab on this!

[stas00]

Awesome! Thank you, @vladdy

Just be aware that the following PR should be merged shortly: #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.

[stas00]

@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!

[stas00]

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.

[vladdy]

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, 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.

[stas00]

@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...

[stas00]

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