Repository: transformers
Issue: Good Second Issue: T5 FP16 in Pytorch
Labels: Good Second Issue

# 🚀 Feature request

This "Good second issue" should revisit some of the problems we were having with FP16 for `T5ForConditionalGeneration`: https://github.com/huggingface/transformers/issues/4586 and help to make T5 compatible with fp16.

**_Requirements:_**

- use transformers master
- use newest pytorch version
- have access to GPU

**_Context:_**

To better explain the context, let's define the three different pre-trained T5 models types we have:

- **T5v1** (original T5): => this corresponds to all those checkpoints: `t5-small`, `t5-base`, `t5-large`, `t5-3b`, `t5-11b`
- **T5v1_1** (improved T5): => this corresponds to all those checkpoints: `google/t5-v1_1-small`, `google/t5-v1_1-base`, `google/t5-v1_1-large`, `google/t5-v1_1-xl`, `google/t5-v1_1-xxl`. **T5v1_1** has a slightly different architecture than **T5v1**. More info on differences can be found here: https://github.com/huggingface/transformers/issues/6285
- **MT5** (multi-lingual T5): => this model is identical in architecture to **T5v1_1** but has different pre-trained weights and a much larger word embedding matrix.

As shown in this issue https://github.com/huggingface/transformers/issues/4586 , training **T5v1** in fp16 mode led in the past to numerical overflow in the `T5LayerFF` forward pass: https://github.com/huggingface/transformers/blob/6189ae99603bd5dc14c5631f1b4562f78e24d575/src/transformers/models/t5/modeling_t5.py#L279. 

At the time of this issue: https://github.com/huggingface/transformers/issues/4586, **T5v1** was added with a small bug that led to slightly wrong outputs that was only fixed by this PR: https://github.com/huggingface/transformers/pull/8518. 

Also, now there are new T5 checkpoints, notably the **T5v1_1** and **MT5** checkpoints, where it would be very interesting to see whether fp16 can work with those.

**_Feature Request_**

So for this feature request, we should two scenarios:

1) Inference:

For each T5 model type we should test when the models break during inference. This can be as easy as testing the following script for a bunch of different checkpoints on different `input_str`:

```python
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch

checkpoint = "t5-small"  # "google/mt5-small", "google/t5-v1_1-small"

model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint).to('cuda')
tokenizer = AutoTokenizer.from_pretrained(checkpoint)

input_str = "Hello there. This is the input."  # here it would be better to test much larger inputs

input_ids = tokenizer(input_str, return_tensors="pt").input_ids.to('cuda')

# FP32
output_fp32 = model.generate(input_ids)

# FP16
model.half()
output_fp16 = model.generate(input_ids)

if output_fp32.tolist() == output_fp16.tolist():
    print("SUCCESS: Output is equal!")
else:
    print("Output is different!")
    print("FP32", output_fp32)
    print("FP16", output_fp16)
```

2) Training (the more interesting part):

This is probably more important and will require more time / skill. In order to check how T5 does in FP16 training, I'd recommend to use the newly added `Seq2SeqTrainer`: https://github.com/huggingface/transformers/blob/6189ae99603bd5dc14c5631f1b4562f78e24d575/src/transformers/trainer_seq2seq.py#L38. I would recommend to train on a summarization task, such as CNN/Dailymail. One could closely follow, this notebook: https://colab.research.google.com/drive/1Ekd5pUeCX7VOrMx94_czTkwNtLN32Uyu?usp=sharing, but replacing Bert2Bert with the different T5 models. Ideally different "fp16 backends" should be tested: https://github.com/huggingface/transformers/blob/6189ae99603bd5dc14c5631f1b4562f78e24d575/src/transformers/training_args.py#L216 and one should try to see whether hacks as proposed in https://github.com/huggingface/transformers/issues/4586#issuecomment-748336815 can solve the problem. It would be very interesting to see whether the error happens only for **T5v1** or also for **T5v1_1** and **MT5** and it what point. For each type it would be great to test for "small", "base" and if possible even "large". Ideally, one should first create a short summarization fine-tuning script (happy to help here) and then run a bunch of different experiments with different fp16 backends and different models. 

**_Possible Outcome_**

The results of those experiments should be documented here or even better on https://discuss.huggingface.co/. Ideally, a solution to the problem is found and one could publish a nice blog post explaining how to effectively train T5.

## Motivation

T5 is one of the most widely used models of Transformers at the moment so that more results to this issue would be extremely useful for the community. In addition, this issue can be a great opportunity to learn more about the limits of fp16 and why some models still do require full fp32 support (or at least until bfloat16 is better supported in torch). This is not an easy issue to tackle, but an extremely important one.

## Your contribution

I'm happy to help along the way, starting with making a nice T5 summarization training pipeline that lets one easily test on different models, and fp16 backends.