

TITLE
[T5/MT5] resolve inf/nan under amp (mixed precision) 

COMMENTS
39

REACTIONS
+1: 2
-1: 0
laugh: 0
hooray: 0
heart: 0
rocket: 0
eyes: 0

BODY
As reported in multiple issues t5/mt5 models produce loss of `nan` under mixed precision training, starting with t5-large and mt5-small and up. This PR is an attempt to fix this issue. This is crucial for DeepSpeed where it's always mixed precision training.

I spent some time with the debugger and the new `detect_overflow` helper util (added in this PR) and discovered that the best place to fix the whole problem is to not `T5LayerFF` in mixed precision. This slightly slows things down/consumes more gpu memory, but no longer requires clamping and running after ever overflowing `hidden_states`.

This PR:

* turns `autocast` off during `T5LayerFF` if run under amp
* removes the previous attempt to clamp the values as it now works without it
* introduces `debug_utils.py` with a helper function `detect_overflow` which is super-handy for tracking overflows automatically (as it's silent if all goes well). It also has some extra features, such as reporting a number of large elements - disabled by default.

Important:
* The fix is only for pytorch built-in amp. apex still has this problem since I haven't researched if the same could be done there, but it's probably a waste of time since apex is being phased out. And deepspeed doesn't use amp so it's till affected.

## Variations

Other possible variations to this solution: 

1. to do the `autocast` disabling dynamically. That is trying with `autocast` and checking if any elements of output are `inf` (not sure of the overhead) and re-running this layer in full fp32 and setting a flag to continue in fp32 from then on. Here the main price will be paid by models that don't need this workaround, but they will gain but not having `autocast` turned off - so it might still be a beneficial solution to all
2. give users a switch to turn this feature on if they discover they need it - or have it on by default and allow users to turn it off if they "know what they are doing".

I am suggesting this since I don't know if all t5/mt5 models are impacted. Definitely t5-small doesn't need this.


## Penalizing large activation

See the details comment: https://github.com/huggingface/transformers/pull/10956#issuecomment-820712267

```
@@ -1578,6 +1618,15 @@ class T5ForConditionalGeneration(T5PreTrainedModel):
             loss = loss_fct(lm_logits.view(-1, lm_logits.size(-1)), labels.view(-1))
             # TODO(thom): Add z_loss https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/layers.py#L666

+            # z_loss
+            log_z = lm_logits.view(-1).logsumexp(-1)
+            z_loss = 7e-5
+            loss_extra = z_loss*log_z.square()
+            #z_loss = 1e-5
+            #loss_extra = z_loss*log_z.pow(3)
+            #print(f"loss={loss}, loss_extra={loss_extra}")
+            loss += loss_extra
```


## Questions:

* If this solution solves the problem at large and is accepted then we probably should document somewhere in t5/mt5 docs that it won't run AMP 100%?

* Test is needed: any suggestions to how we could write a test that is not too big and still gets nans prior to this PR? `t5-small` and `t5-base` don't have this problem (at least with a small sample), in my experiments the first model that gets `inf/nan` on the first batch is `mt5-small` (1.2GB), so my minimal test is:
```
rm -rf output_dir; CUDA_VISIBLE_DEVICES=0 USE_TF=0 PYTHONPATH=src python examples/seq2seq/run_translation.py \
--model_name_or_path google/mt5-small --do_train  --source_lang en --target_lang ro --dataset_name wmt16 \
--dataset_config_name ro-en --output_dir output_dir --per_device_train_batch_size=4  --logging_step 2 --save_steps 0 \
--fp16 --max_train_samples 10 --save_total_limit 0 --save_strategy no
```
We can then run this as a test and check for `nan` in loss reports. 

But the 1.2GB download is somewhat big even for `@slow` tests. 
**edit**: @LysandreJik says it's not a problem since we are now caching the models on the test machine.

If it is ok I will just stick this with all the extended tests under `examples/tests/trainer/test_trainer_ext.py` where we have a setup for this type of full application-based tests.

* I also know some users mentioned that `inf` may happen much later in the game. I haven't run very long tests.

TODO:

* [ ] I left all the debug prints in place so that you could experiment with it easily - will remove when this is approved to be a good change

Related discussions:

- https://discuss.pytorch.org/t/bfloat16-transformers/96260 pegasus is affected too

Fixes: https://github.com/huggingface/transformers/issues/10830
Fixes: https://github.com/huggingface/transformers/issues/10819

@patrickvonplaten, @patil-suraj, @LysandreJik 