[T5/MT5] resolve inf/nan under amp (mixed precision) 
Author: stas00
Pull request: yes
Labels: name: WIP

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.

---------

FYI, there is another solution posted by @ibeltagy here: https://github.com/huggingface/transformers/issues/14189#issuecomment-961571628 it is based on custom scaling. 

--------

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 

Comments:
> But the 1.2GB download is somewhat big even for @slow tests.

The downloads are cached on a shared disk across slow self-hosted runners, so that's not an issue!

Before I approached this problem, I did a bit of a study on the bfloat16 vs float16 properties. This is not fully complete, but you can see most of the useful data here: https://github.com/stas00/ml-ways/blob/master/numbers/bfloat16-vs-float16-study.ipynb

Comments/requests/suggestions are welcome though. It's a bit on a terse-side.

I spent some more time staring at the numbers, as I think @patrickvonplaten mentioned in one of the related threads, something trained in `bfloat16` isn't going to work with `float16`. You can see why by looking at this debug output:

```
min=-2.77e+05 max= 2.75e+05 var= 5.45e+07 mean= 5.16e+01 (T5Stack loop start)
min=-2.77e+05 max= 2.75e+05 var= 5.45e+07 mean= 5.16e+01 (T5Block)
min=-2.77e+05 max= 2.75e+05 var= 5.45e+07 mean= 5.16e+01 (T5LayerNorm)
min= 1.31e+06 max= 6.90e+08 var= 9.52e+15 mean= 5.45e+07 (T5LayerNorm variance)
min=-1.46e+01 max= 1.46e+01 var= 1.00e+00 mean=-2.69e-03 (T5LayerNorm hidden_states)
min=-1.46e+01 max= 1.46e+01 var= 1.00e+00 mean=-2.69e-03 (T5LayerNorm hidden_states before return)
min=-2.76e+05 max= 2.74e+05 var= 5.41e+07 mean= 4.83e+01 (T5Block after T5LayerSelfAttention)
min=-2.76e+05 max= 2.74e+05 var= 5.41e+07 mean= 4.83e+01 (T5LayerNorm)
min= 1.38e+06 max= 6.86e+08 var= 9.37e+15 mean= 5.41e+07 (T5LayerNorm variance)
min=-1.45e+01 max= 1.46e+01 var= 1.00e+00 mean=-2.98e-03 (T5LayerNorm hidden_states)
min=-1.45e+01 max= 1.46e+01 var= 1.00e+00 mean=-2.98e-03 (T5LayerNorm hidden_states before return)
min=-2.76e+05 max= 2.73e+05 var= 5.40e+07 mean= 3.93e+01 (T5Block before T5LayerFF)
min=-2.76e+05 max= 2.73e+05 var= 5.40e+07 mean= 3.93e+01 (T5LayerFF: 1)
min=-2.76e+05 max= 2.73e+05 var= 5.40e+07 mean= 3.93e+01 (T5LayerNorm)
min= 1.61e+06 max= 6.84e+08 var= 9.28e+15 mean= 5.40e+07 (T5LayerNorm variance)
min=-1.44e+01 max= 1.46e+01 var= 1.00e+00 mean=-5.14e-03 (T5LayerNorm hidden_states)
min=-1.44e+01 max= 1.46e+01 var= 1.00e+00 mean=-5.14e-03 (T5LayerNorm hidden_states before return)
min=-2.47e+00 max= 3.03e+00 var= 4.43e-02 mean=-8.23e-05 (T5LayerFF: 2)
min=-1.70e-01 max= 4.95e+01 var= 6.34e-01 mean= 3.00e-01 (gelu 1)
min=-3.70e+02 max= 3.93e+02 var= 3.79e+02 mean= 2.79e-01 (gelu 2)
min=-4.71e+03 max= 3.67e+03 var= 1.89e+03 mean=-3.80e-01 (gelu 3)
min=-5.23e+03 max= 4.08e+03 var= 2.21e+03 mean=-4.75e-01 (gelu 4)
min=-7.11e+04 max= 5.32e+04 var= 8.27e+06 mean=-1.36e+02 (gelu 5)
min=-7.11e+04 max= 5.32e+04 var= 8.27e+06 mean=-1.36e+02 (T5LayerFF: 3)
min=-2.61e+05 max= 2.68e+05 var= 4.41e+07 mean=-1.04e+02 (T5LayerFF: 5)
min=-2.61e+05 max= 2.68e+05 var= 4.41e+07 mean=-1.04e+02 (T5Block after T5LayerFF)
min=-2.61e+05 max= 2.68e+05 var= 4.41e+07 mean=-1.04e+02 (T5Stack loop end)
min=-2.61e+05 max= 2.68e+05 var= 4.41e+07 mean=-1.04e+02 (T5LayerNorm)
min= 2.99e+06 max= 6.12e+08 var= 5.65e+15 mean= 4.41e+07 (T5LayerNorm variance)
min=-1.45e+01 max= 1.62e+01 var= 1.00e+00 mean=-2.27e-02 (T5LayerNorm hidden_states)
min=-1.45e+01 max= 1.62e+01 var= 1.00e+00 mean=-2.27e-02 (T5LayerNorm hidden_states before return)
```

Because `bfloat16` lacks precision - it trained itself to compensate for this by switching to the range of large numbers. If you look at the numbers above you can see that many of them are a way beyond fp16 range, which can only do `+-64K`. 

So if I understand the nature of this problem correctly expecting this to work is a bit of fantasy. But of course, let's try to do our best to come as close to the solution as possible.

I found that it's enough to cancel autocast just for `self.DenseReluDense` for the simple case to not produce NaN. 

@yuvalkirstain, let's switch the discussion to the actual PR

wrt your newly discovered overflow.

Please try to add this penalizing for large logits:
```
@@ -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
```

May need some tuning for `z_loss` factor for best convergence. The recommended one is 1e-4, so I've experimented with a few. Also tried the `pow(3)` instead of `pow(2)`.

It seem that the network gets the hint within just a 100 steps - `loss_extra` drops down very quickly.

Perhaps this was the missing piece?


Here is the output of the proposed [overflow/underflow detector](https://github.com/huggingface/transformers/pull/11274) in progress tool for mt5. This is prior to any modifications proposed in this PR. So one can see the progression as the weights and activations change from forward to forward.

```
rm -rf output_dir; CUDA_VISIBLE_DEVICES=0 USE_TF=0 PYTHONPATH=src \
python examples/pytorch/translation/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 --debug underflow_overflow
```

```
Detected inf/nan during batch_number=0
Last 21 forward frames:
abs min  abs max  metadata
                  encoder.block.1.layer.1.DenseReluDense.dropout Dropout
0.00e+00 2.57e+02 input[0]
0.00e+00 2.85e+02 output
                  encoder.block.1.layer.1.DenseReluDense.wo Linear
4.80e-06 8.62e+00 weight
0.00e+00 2.85e+02 input[0]
8.50e-05 1.53e+03 output
                  encoder.block.1.layer.1.DenseReluDense T5DenseGatedGeluDense
0.00e+00 2.04e+00 input[0]
8.50e-05 1.53e+03 output
                  encoder.block.1.layer.1.dropout Dropout
8.50e-05 1.53e+03 input[0]
0.00e+00 1.70e+03 output
                  encoder.block.1.layer.1 T5LayerFF
0.00e+00 1.50e+03 input[0]
6.78e-04 3.15e+03 output
                  encoder.block.1 T5Block
0.00e+00 1.40e+03 input[0]
6.78e-04 3.15e+03 output[0]
             None output[1]
2.25e-01 1.00e+04 output[2]
                  encoder.block.2.layer.0.layer_norm T5LayerNorm
6.54e-02 2.75e-01 weight
6.78e-04 3.15e+03 input[0]
5.75e-06 2.12e+00 output
                  encoder.block.2.layer.0.SelfAttention.q Linear
3.75e-08 3.40e-01 weight
5.75e-06 2.12e+00 input[0]
2.21e-06 1.20e+00 output
                  encoder.block.2.layer.0.SelfAttention.k Linear
4.84e-08 2.62e+00 weight
5.75e-06 2.12e+00 input[0]
5.47e-05 1.40e+01 output
                  encoder.block.2.layer.0.SelfAttention.v Linear
7.21e-06 2.59e+00 weight
5.75e-06 2.12e+00 input[0]
1.20e-04 7.56e+00 output
                  encoder.block.2.layer.0.SelfAttention.o Linear
6.65e-06 1.44e+01 weight
0.00e+00 5.30e+00 input[0]
5.20e-04 2.66e+02 output
                  encoder.block.2.layer.0.SelfAttention T5Attention
5.75e-06 2.12e+00 input[0]
5.20e-04 2.66e+02 output[0]
             None output[1]
2.25e-01 1.00e+04 output[2]
                  encoder.block.2.layer.0.dropout Dropout
5.20e-04 2.66e+02 input[0]
0.00e+00 2.96e+02 output
                  encoder.block.2.layer.0 T5LayerSelfAttention
6.78e-04 3.15e+03 input[0]
2.65e-04 3.42e+03 output[0]
             None output[1]
2.25e-01 1.00e+04 output[2]
                  encoder.block.2.layer.1.layer_norm T5LayerNorm
8.69e-02 4.18e-01 weight
2.65e-04 3.42e+03 input[0]
1.79e-06 4.65e+00 output
                  encoder.block.2.layer.1.DenseReluDense.wi_0 Linear
2.17e-07 4.50e+00 weight
1.79e-06 4.65e+00 input[0]
2.68e-06 3.70e+01 output
                  encoder.block.2.layer.1.DenseReluDense.wi_1 Linear
8.08e-07 2.66e+01 weight
1.79e-06 4.65e+00 input[0]
1.27e-04 2.37e+02 output
                  encoder.block.2.layer.1.DenseReluDense.dropout Dropout
0.00e+00 8.76e+03 input[0]
0.00e+00 9.74e+03 output
                  encoder.block.2.layer.1.DenseReluDense.wo Linear
1.01e-06 6.44e+00 weight
0.00e+00 9.74e+03 input[0]
3.18e-04 6.27e+04 output
                  encoder.block.2.layer.1.DenseReluDense T5DenseGatedGeluDense
1.79e-06 4.65e+00 input[0]
3.18e-04 6.27e+04 output
                  encoder.block.2.layer.1.dropout Dropout
3.18e-04 6.27e+04 input[0]
0.00e+00      inf output
```

Hi there, I'm wondering what the current status of this is, as my team would benefit from a fix to fp16 issue with large T5 models. And is there anything we could do to help to move the PR along? 

In the mean time, it should be sufficient to simply disable autocast for the DenseReluDense, correct?

> Hi there, I'm wondering what the current status of this is, as my team would benefit from a fix to fp16 issue with large T5 models. And is there anything we could do to help to move the PR along?

@yuvalkirstain, who is one of the original reporters mentioned elsewhere that he still had an issue during the long training, so I was waiting for him to provide more details.

> In the mean time, it should be sufficient to simply disable autocast for the DenseReluDense, correct?

If you're not using deepspeed, then yes, that is all that is needed. At least for the tests I have done. But they weren't long.

Perhaps you could test this PR and report back if it solves your problem?

I'm not sure if I should remove the clamping or not.

I cleaned up the PR to remove all the debug noise, so it's very simple now.

Hi, I ran some experiments and it appears to me that this branch does fix the inf/nan issue for both T5-large and T5-3b--I trained both models for 10,000 steps on a language modeling task and never had the NaN loss issue I was having before. However, as far as I can tell the fix comes at a large cost in time and memory usage. 

Using t5-large on an A6000 card (48 GB), I found:

- no fp16: 25.00 GB, 3.06 iters/s
- fp16 without the fix from this branch: 15.01 GB, 4.10 iters/s [but loss was `NaN`]
- fp16 with the fix from this branch: 23.99 GB, 2.90 iters/s

(collected using the `torch.autograd.profiler` tool)

In other words, fp16 with this fix uses about 1.6x more memory than before.

Disclaimer: the experiments I ran were using an LM task that's internal to my team, so you won't be able to replicate it exactly. But I wanted to report back anyway since it's been a few days. However, in the next few days I'd like to repeat these experiments using one of the HF example scripts so that you can verify by running the exact same code.

That's a fantastic feedback, @dblakely - Thank you! Looking forward to seeing the stats on non-custom code.

It's interesting that you get even slower results than full fp32. But since you're no A6000 you're probably running on tf32 automatically if you're on the recent pytorch, that would explain it.

And I trust you're not using the overflow detector which would add to the slowdown a bit.

BTW, apparently there is a new `torch.profiler` tool - I haven't tried it yet.

Also earlier I wrote:

> I found that it's enough to cancel autocast just for `self.DenseReluDense` for the simple case to not produce NaN.

So you might want to try this slightly tighter version:

```
class T5LayerFF(nn.Module):
    def __init__(self, config):
        super().__init__()
        if config.feed_forward_proj == "relu":
            self.DenseReluDense = T5DenseReluDense(config)
        elif config.feed_forward_proj == "gated-gelu":
            self.DenseReluDense = T5DenseGatedGeluDense(config)
        else:
            raise ValueError(
                f"{self.config.feed_forward_proj} is not supported. Choose between `relu` and `gated-gelu`"
            )

        self.layer_norm = T5LayerNorm(config.d_model, eps=config.layer_norm_epsilon)
        self.dropout = nn.Dropout(config.dropout_rate)

    def forward(self, hidden_states):
        forwarded_states = self.layer_norm(hidden_states)
        if torch.is_autocast_enabled():
            with torch.cuda.amp.autocast(enabled=False):
                forwarded_states = self.DenseReluDense(forwarded_states)
        else:
            forwarded_states = self.DenseReluDense(forwarded_states)
        hidden_states = hidden_states + self.dropout(forwarded_states)
        return hidden_states
```

but given that the bulk of everything comes from `DenseReluDense` it probably won't make much of a difference speed and memory requirements-wise.


Hi all,

If `bf16` is what's native to these models, how about we do `autocast` with `bf16` instead of `fp16` (and then don't scale)?  There is a pull request [here](https://github.com/pytorch/pytorch/issues/55374) to add a `bf16` option to autocast.

That would be the best solution assuming you have high-end Ampere GPUs which support bf16 natively. (rtx-3090, a100, ...). So once this is finalized in pytorch we will support it in the HF trainer as well. 

If you have been actively watching that development please kindly ping us when it's completed in pytorch. Thank you.

We use rtx a6000s, so I believe we are ok on that front.  I'll monitor the aforementioned PR and keep you updated

UPDATE: the fix has migrated to [this pr](https://github.com/pytorch/pytorch/pull/61002)

The [torch pr](https://github.com/pytorch/pytorch/pull/61002) is almost through, so I'm coming back to this.  Would the ensuing pr here be as simple as changing autocast in t5 to the bf16 option?

Thank you for keeping on top of torch's side, @JamesDeAntonis 

No, we will have to rework the HF Trainer to support bf16. The model doesn't need to be changed.

My recommendation is to wait till that PR lands in pt-nightly so we have something to test with. And then we can work on having bf16 support in the trainer.

If you're not using the HF Trainer, then you can do it independently by wrapping the training step in the new autocast.

Sounds good, thanks for the quick response!  I'll continue to watch the pr.

We do indeed use the HF Trainer, so I'll probably be active on the HF pr as well.

It looks like the PR was just merged in torch!  I think the ball is now in our court once the nightly build hits (so I think starting tomorrow)

Awesome! Thank you for keeping us abreast of this development, @JamesDeAntonis.

This is a month of August and most team members are on vacation at the moment, so this might take longer than normal.

My plate is very full at the moment, so unless someone beats me to it, I probably won't have any time in the next few weeks to work on this.

But, first, please create a new Issue and tag me there, so that we have an easy way to track this feature request.

Second, if one of you would like to work on the PR to integrate bf16 that would be great. I think the change itself should be relatively simple, add a new CLI arg `--bf16` and set amp to bf16 instead of fp16 in trainer.py. We may have to deprecate `fp16_backend` and rename it to something more generic, but just doing the above is a good start. The devil is in the detail though, so it may take longer to figure out.

Hi, after reading all the replies and related issues in torch and transformers, I still don't know how to fix the nan problem. I get `loss=nan` on every simple example using mt5-base.

```
from transformers import AutoConfig, AutoTokenizer, AutoModelForSeq2SeqLM, AutoModel
import torch
device = 'cuda'
tokenizer = AutoTokenizer.from_pretrained('google/mt5-base')
model = AutoModelForSeq2SeqLM.from_pretrained('google/mt5-base')
model.to(device)
scaler = torch.cuda.amp.GradScaler()
toks = tokenizer(['Je vous invite à vous lever pour cette minute de silence.',"Please rise, then, for this minute' s silence."], return_tensors='pt', padding='max_length',max_length=512, truncation=True).to(device)
inputs = {'input_ids': toks['input_ids'][0:1], 'attention_mask': toks['attention_mask'][0:1], 'labels': toks['input_ids'][1:], 'output_hidden_states':True}
with torch.cuda.amp.autocast():
    outputs = model(**inputs)
    loss = outputs.loss
scaler.scale(loss).backward()
loss.item()

> nan

```
torch=1.8.1+cu101
transformers=4.9.2
any help would be much appreciated!

>Hi, after reading all the replies and related issues in torch and transformers, I still don't know how to fix the nan problem. I get `loss=nan` on every simple example using mt5-base.

If you want to use `bf16`, you need to include the `fast_dtype` as seen [here](https://github.com/pytorch/pytorch/blob/master/torch/autocast_mode.py#L128) (I think)

> > Hi, after reading all the replies and related issues in torch and transformers, I still don't know how to fix the nan problem. I get `loss=nan` on every simple example using mt5-base.
> 
> If you want to use `bf16`, you need to include the `fast_dtype` as seen [here](https://github.com/pytorch/pytorch/blob/master/torch/autocast_mode.py#L128) (I think)

Thank you very much. I switch to pt nightly build 1.10.0+cu111 and test the `fast_dtype =torch.bfloat16` in `amp`.  It seems that cuda device does not support bfloat16. 

```
RuntimeError: Current CUDA Device does not support bfloat16. Switching fast_dtype to float16.
```
I tried Tesla M40, GTX TITAN X  and Quadro RTX 8000, same error.


> I tried Tesla M40, GTX TITAN X and Quadro RTX 8000, same error.

For bf16 you want the high end Ampere cards https://en.wikipedia.org/wiki/Ampere_(microarchitecture)#Products_using_Ampere - so 3090 and below on that list.

this PR isn't merged yet? :( is the issue resolved?

Well, it introduces a small slowdown as it forces one FF in fp32 under mixed precision, so I wasn't sure whether this solves the problem for everybody. Or whether this should be configurable. Some users reported that it solved their problem, other that it didn't.

Additionally I proposed in this PR 
- ~to remove the clamping, but got no feedback whether it's safe to do. https://github.com/huggingface/transformers/pull/10956/files/1ddec2c860617230a5171f3a95be74d27f4c8e9d#r603663224~ (Patrick suggested to keep it, so I restored that part)
- to add an additional penalty to lm loss as described in the original codebase (it's not in PR, but the code is in the OP), which would stir the finetuning into the direction of fp16 weights. Perhaps it should be only added for when autocast is detected? But then bf16 is imminent, so probably need to find a way to check the autocast dtype is fp16? (and the dtype was introduced in pt-1.10 only)

Here is a possible plan of action:

1. leave everything as is and just have this PR add the FF override in fp32
2. ~discuss clamping and keep or remove it~ (Patrick suggested to keep it, so I restored that part)
3. discuss large weight lm loss penalty factor and add the code in


This PR forces T5 FF Layer in fp32. With this change, there is almost no benefit to training in fp16.
The memory usage and training speed improvements are very limited. 

@Liangtaiwan, also you may to try the loss penalty factor. The patch to apply (instead of this PR) is in the OP.

Hi everyone, I am testing a method of adjusting the T5 weights for FP16 training and so far it's promising. However, I would like to see if there is a way to "validate" how much of the model performance is still retained for both pre-trained and fine-tuned tasks.

The TLDR is: we scale the weights down, for as few parameters and as little as possible, until the model can be trained without NaN. Basically, to perform the minimum amount of "surgery" on the weights.

Currently, I am reducing about 2-3% of parameters in the model by a factor of 2 only and seeing some good initial results. These parameters are in the feed-forward layers in the encoder only. The resulting model still seems to work on existing tasks and I can fine-tune T5-large just fine in FP16 on my own task, where previously it would NaN. So far, nothing seems to be wrong with the outputs, and I have not encountered NaN. 

For example of a converted model: https://github.com/tlkh/t5-fp16-surgery/blob/main/t5-large.ipynb

I have uploaded the converted models for people to play with: 

* [`tlkh/t5_large_fp16_untuned`](https://huggingface.co/tlkh/t5_large_fp16_untuned)
* [`tlkh/t5_3B_fp16_untuned`](https://huggingface.co/tlkh/t5_3B_fp16_untuned)

Note: for the 3B model, after the conversion, the pre-trained translation task seems to be more unstable, but given it can still generate coherent text, my hunch is that after fine-tuning on another task, it should have negligible difference. However, it would be great to know for sure, so wonder if there is some kind of benchmark suite we can try.

I do not have the resources to convert the 11B one, but I do not see why that wouldn't work similarly. It is also very quick to convert models. 

GitHub repo to demo/show code for the conversion, and also included inference testing to show the model seems to be working fine: https://github.com/tlkh/t5-fp16-surgery

My hopeful outcome from this is that we can fine-tune T5 in FP16 without any real penalty. 

I'm trying to use this for T5-3B with A100. Is bf16 available experimentally? 

BTW, I'm heading towards this direction because the `fairscale sharded_ddp` option for some reason hangs when it should run evaluation. Any pointers to solve this issue as well?

> I'm trying to use this for T5-3B with A100. Is bf16 available experimentally?

There is a WIP PR: https://github.com/huggingface/transformers/pull/13207

> BTW, I'm heading towards this direction because the `fairscale sharded_ddp` option for some reason hangs when it should run evaluation. Any pointers to solve this issue as well?

Use deepspeed: https://huggingface.co/transformers/master/main_classes/deepspeed.html#deepspeed-trainer-integration


> @Liangtaiwan, also you may to try the loss penalty factor. The patch to apply (instead of this PR) is in the OP.

@stas00 Could you point out where is the patch or the PR? 



https://github.com/huggingface/transformers/pull/10956 Scroll down to "Penalizing large activation"

FYI, there is another solution posted by @ibeltagy here: https://github.com/huggingface/transformers/issues/14189#issuecomment-961571628

it too is based on custom scaling. Please have a look.

cc: @tlkh 


So heads up to all those watching this PR - if you have Ampere GPUs start using `--bf16` which was just added (i.e. use `master`) and the overflow problem will be no more: https://huggingface.co/docs/transformers/master/performance#bf16

p.s. I haven't actually tested that it's so with this particular issue, so if you do and find that it is not so, please kindly flag this to me.

So I got a chance to test the new `--bf16` flag on this issue with RTX-3090 (Ampere) and now mt5 doesn't overflow:

```
rm -rf output_dir; CUDA_VISIBLE_DEVICES=0 USE_TF=0 PYTHONPATH=src python \
examples/pytorch/translation/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 --max_train_samples 10 \
--save_total_limit 0 --save_strategy no --bf16

***** train metrics *****
  epoch                    =        3.0
  train_loss               =    28.7758
  train_runtime            = 0:00:01.94
  train_samples            =         10
  train_samples_per_second =     15.458
  train_steps_per_second   =      4.637
```

we get the same outcome with fp32 (i.e. w/o  `--bf16`). 

With `--fp16` we still overflow (no surprise here, I have just re-checked that):

```
rm -rf output_dir; CUDA_VISIBLE_DEVICES=0 USE_TF=0 PYTHONPATH=src python \
examples/pytorch/translation/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 --max_train_samples 10 \
--save_total_limit 0 --save_strategy no --fp16

***** train metrics *****
  epoch                    =        3.0
  train_loss               =        0.0
  train_runtime            = 0:00:01.74
  train_samples            =         10
  train_samples_per_second =      17.24
  train_steps_per_second   =      5.172
```

> This PR forces T5 FF Layer in fp32. With this change, there is almost no benefit to training in fp16. The memory usage and training speed improvements are very limited.

Hi, same issue here. Did you find a way to fix it?

@GaryYufei 
I don't think it's almost impossible to fix this issue.
The best way is to use GPUs that support bf16 training. 
You may also try to use @tlkh proposed method but not sure there would be any side-effect.

In my case, I choose to use fp32 to finetune the model. 




The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_10956). All of your documentation changes will be reflected on that endpoint.

Is this problem solved in the latest release of transformers?

It's not really a problem in `transformers` per se, but a limitation of the model.

This PR didn't get merged as it helped only in some cases and it would introduce a slowdown.

This thread contains various possible workarounds, but the best solution at the moment is to use bf16-able hardware to finetune t5 and any other bf16-pretrained models (Ampere GPUs or TPUs)

I see, thank you for your answer

Can someone approve this ? I'm getting nan values .. on main branch

It hasn't been merged because it's not the ideal solution as it introduces a degradation in performance (scales to fp32 = more memory used) and it doesn't work always resolve the problem.

This is a curse of many bf16-pre-trained models used in fp16 mode and not just of T5 and its derivatives. 

Do you by chance have access to Ampere gpus and are able to use bf16 instead of fp16 - this would solve the problem w/o changing the code. https://github.com/huggingface/transformers/pull/10956#issuecomment-997230177


In my case, the issue with t5 training, causing nan values in mix precision (using torch native amp), has been resolved following this PR. Thank you very much. Have a great day!

Closing as this PR is super old and not planned