FP-16 training producing nans on t5-large/flan-t5-xl
Author: cassianlewis
Pull request: no

### System Info

This was an issue a while back but seems to have resurfaced - https://discuss.huggingface.co/t/t5-fp16-issue-is-fixed/3139

I have tested the exact following code on `t5-small` and `t5-base` and they work fine. However, when using `t5-large` and/or `flan-t5-xl`, the model produces nan outputs. This is solely a result of using half precision (ignore the multiple GPUs, strategy etc, I have tested with every other variation):

```
trainer = pl.Trainer(
    precision="16",
    accelerator='gpu',
    strategy='auto',
    devices=4,)
```

I am using `transformers == 4.28.1` and `lightning == 2.0.0`

Any ideas/help appreciated
Thanks!


### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

```
trainer = pl.Trainer(
    precision="16",
    accelerator='gpu',
    strategy='auto',
    devices=4,)
```



### Expected behavior

Nans!!!

Comments:
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

Hi @cassianlewis, thanks for reporting this issue. 

There's been a few updates recently regarding handling of loading of model weights and their precision. Could you update your transformers version to the latest release and let us know if the issue still persists? 

Could you also provide some more details about the other strategies tested (when this behaviour doesn't occur)? 

cc @younesbelkada 

Hi @amyeroberts, I think the issue is more with the model itself. See https://www.philschmid.de/fine-tune-flan-t5-deepspeed

`When fine-tuning T5 models we cannot use fp16 since it leads to overflow issues, see: [#4586](https://github.com/huggingface/transformers/issues/4586), [#10830](https://github.com/huggingface/transformers/issues/10830), [#10956](https://github.com/huggingface/transformers/pull/10956)`

(t5 was trained with bf16)

FYI this behaviour doesn't occur when I remove the `precision="16"` line 

@cassianlewis Thanks for providing these details and links. 

Looking at the linked issues and Philip's blog, unfortunately I don't think there's much that can be done to resolve this. We can leave this issue open for now for people to share any updates or solutions they may have. 

Hi @cassianlewis 
I think it is expected that pure fp16 training leads to unstable behaviour for some models and in some configuration; have you tried to train the models in bf16 precision?

@younesbelkada My Tesla T4 doesn't currently support bf16 (you need Ampere architecture), so I will get back to you on that 

I was running into the same issue, and I found a few things that helped (but did make training time much slower). 

1. Specifically, I found that I needed to add the line `@torch.cuda.amp.custom_fwd(cast_inputs=torch.float32)` to the forward method of this module https://github.com/huggingface/transformers/blob/2ab75add4b30c2fc44a8bf575156d448d9ed87a7/src/transformers/models/t5/modeling_t5.py#L301-L314

Of course, this greatly slows things down since `fp16` won't be utilized at all in this module anymore (and this would slow down other methods like bf16 training). I think if there is a way to enforce that `hidden_linear` should always stay in `fp32` then the `torch.cuda.amp...` statement would not be needed, but I was unable to figure out how to get that to work.

2. Another thing I changed were these blocks in the code https://github.com/huggingface/transformers/blob/2ab75add4b30c2fc44a8bf575156d448d9ed87a7/src/transformers/models/t5/modeling_t5.py#L707-L713
to use the alternate if statement 
```python
if torch.isinf(hidden_states).any():
    ...
```
since I found it was not guaranteed that `hidden_states` would be in `torch.float16` during fp16 training, but infinity values could still be present. 

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.