Pegasus pretraining in fp16 results in NaN loss
Author: kolakows
Pull request: no

## Environment info

`transformers` version: 4.5.1
- Platform: Linux-5.4.0-73-generic-x86_64-with-glibc2.29
- Python version: 3.8.5
- PyTorch version (GPU?): 1.4.0 (True)
- Tensorflow version (GPU?): not installed (NA)
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: no

### Who can help
@patrickvonplaten, @patil-suraj

## Information

Model I am using: pegasus

The problem arises when using:
* [ ] my own modified scripts:
```
config = PegasusConfig(**pegasus_config_kwargs)
model = PegasusForConditionalGeneration(config=config)
```
and then using Trainer with fp16 on.

The trainer args I'm using:
```json
{
  "logging_strategy": "steps",
  "logging_steps": 20,
  "save_strategy": "steps",
  "save_steps": 5000,
  "num_train_epochs": 2,
  "lr_scheduler_type": "linear",
  "warmup_steps": 10000,
  "learning_rate": 0.001,
  "dataloader_num_workers": 8,
  "per_device_train_batch_size": 16,
  "gradient_accumulation_steps": 16,
  "group_by_length": true,
  "adafactor": true,
  "fp16": true
}
```

The tasks I am working on is:
* [ ] my own task or dataset

## To reproduce

I was trying to pretrain pegasus in fp16 from scratch using a modified script. The training is much faster, around 40% speedup, but after almost 3 days, training was 10% into a second epoch, a NaN loss happened.  Debugging the place where overflow occurred I guess is possible, but will be troublesome. Do you know what could be the problem or if someone is working on problems with fp16 on pegasus?

I've seen for example that it could be a problem when using pretrained checkpoints (https://discuss.huggingface.co/t/finetuning-for-fp16-compatibility/977), but shouldn't it work when initializing model from config, like below?
```
config = PegasusConfig(**pegasus_config_kwargs)
model = PegasusForConditionalGeneration(config=config)
```

Training without fp16 works fine.

Comments:
Yeah, not sure to what extent it is feasible to prevent this as Pegasus was pretrained in `bfloat16` cc @stas00 

But I'm pretraining a freshly initialized model, so I think the problem shouldn't be with the `bfloat16` casting

That's interesting. We have primarily debugged bf16-pretrained models that almost all had this issue as Patrick says.

So this means the model's design is somehow not fp16-friendly.

Could you take a last checkpoint that was still good and run it with `DebugUnderflowOverflow`
https://huggingface.co/transformers/debugging.html#underflow-and-overflow-detection
and report back the failing trace - which will show us where the under/over-flow occurs in the model.


I will debug it, thanks for the link on how to do it, but probably will have the results in like ~2 weeks time, because now I'm waiting for the results of training without mixed precision. 

I've run training from checkpoint with debugging, like below:

```
DebugUnderflowOverflow(model)
trainer = Trainer(**args)
trainer.train(resume_from_checkpoint=checkpoint_path)
```

And NaN's happened after around 2 days of training. I've redirected all stdout to file,  but the problem is that there wasn't any output from the DebugUnderflowOverflow, no near first nans or other place in file. Console also didn't show anything.

```
Logging step 78200
{'loss': 3.6594, 'learning_rate': 0.00028089413749472796, 'epoch': 0.54}
Logging step 78220
{'loss': nan, 'learning_rate': 0.0002806832560101223, 'epoch': 0.54}
```
Assuming that I've used DebugUnderflowOverflow correctly, do you have any ideas what might be the source of these nans?

Disclaimer about the experiment: 
Last checkpoint I had was just before the end of the first epoch, the next one was after NaN's started, so I took the last one, but because we have 'dynamic' tokenization it would take long time just to get to the previous point. So I've used option ignore_data_skip which didn't forward the dataset to the checkpoint's data point but just started training on the whole dataset again. I think it shouldn't matter for the purpose of debugging NaNs because in the first run model has seen whole training dataset without throwing NaN's.

You can validating that the tracing works with:
https://huggingface.co/transformers/debugging.html#specific-batch-absolute-mix-and-max-value-tracing

This will just report all min/max values of the desired batches - e.g. batch 0, so that you know it's configured correctly and outputs the data it would if there were to be NaNs.

Let's validate that it works first and if it does, then hopefully a trace of one batch could shed some light. If it's really long probably make an attachment to your comment.

e.g. it's possible that the weights are all not-NaNs, but the loss still somehow gets pushed into over/underflow.

Turns out I didn't attach the debugger properly for the first time ^^. A way to validate helped, thanks.

Here are all frames printed by the debugger after detecting inf.
[overflow_debug.txt](https://github.com/huggingface/transformers/files/6777740/overflow_debug.txt)

Not sure where things go wrong.

The last frame is:
```
                  model.encoder.layers.10 PegasusEncoderLayer
2.44e-04 6.86e+04 input[0]
0.00e+00 3.40e+38 input[1]
9.77e-04      inf output[0]
```
The weird thing is that, input[1] is probably attention_mask and not sure why and where some of its values are set to inf. I think in the encoder layer it should be 0 or 1, indicating padding masking? 

Other thing I don't fully understand is that

```
                  model.encoder.layers.10.fc2 Linear
1.52e-08 3.28e+00 weight
3.00e-04 1.71e+00 bias
0.00e+00 1.80e+02 input[0]
0.00e+00 6.06e+04 output
                  model.encoder.layers.10 PegasusEncoderLayer
2.44e-04 6.86e+04 input[0]
0.00e+00 3.40e+38 input[1]
9.77e-04      inf output[0]
```

it looks like second layer of feed forward layer returns output that is still acceptable in fp16, but then the whole layer returns inf. So I assume that the overflow occurred somewhere between this line (4.5.1 version) https://github.com/huggingface/transformers/blob/4bae96ec2bee265f938fc262201538819419089a/src/transformers/models/pegasus/modeling_pegasus.py#L337 and the return. 

But there is a check to clamp any possible overflows.

I'd say the next step is to inject the `detect_overflow` between the suspect lines of code, as shown at the very end of:
https://huggingface.co/transformers/debugging.html#underflow-and-overflow-detection

as shown in this example:
```
from debug_utils import detect_overflow

class T5LayerFF(nn.Module):
    [...]
    def forward(self, hidden_states):
        forwarded_states = self.layer_norm(hidden_states)
        detect_overflow(forwarded_states, "after layer_norm")
        forwarded_states = self.DenseReluDense(forwarded_states)
        detect_overflow(forwarded_states, "after DenseReluDense")
        return hidden_states + self.dropout(forwarded_states)
```

and then you will know exactly where things overflow.

And once identified you can either turn off the `autocast` off for that line of code, or to change the operation to always cast to fp32, as in `some_torch_op(...., dtype=torch.float32)` if it's a torch op that is.

For `autocast` turning off example please see https://github.com/huggingface/transformers/pull/10956/files

So I ran some more tests with `detect_overflow`.

Turns out that scaling up inside `F.dropout` pushes already high values from output of 2nd linear layer (which is fp16) into inf. 

The next unexpected thing is that, the inf check
https://github.com/huggingface/transformers/blob/4bae96ec2bee265f938fc262201538819419089a/src/transformers/models/pegasus/modeling_pegasus.py#L340
should be moved a few lines of code up. 

As it is now, during the dtype check, `hidden_states` is already promoted to fp32 after the residual add. In residual add `residual` is fp32 and `hidden_states` is fp16 with possible overflows that get carried to fp32 result.

Moving the check up will patch the overflows I've been seeing. 

I also think about adding a second check before the first residual add in the encoder, as some values are rather high (2.2e4). And then I'll keep my fingers crossed that nothing overflows in the decoder as I haven't looked into the scale of values there. 


I'm glad to hear that you can now easily tell where things overflow, @kolakows.

Please remember that the original code was trained in a different dtype regime (bf16 or fp32/tf32) and so the designers of the model haven't had to deal with fp16 and that's why changes are needed to be applied to the original port. This same story happens to pretty much all models of this kind (i.e. not designed to be trained with fp16 in mind).

I trust you will be able to tweak the code to overcome this.

You can approach this in 3 ways 
1. explicit upcasting as you suggested
2. turning off `autocast` for the duration of the "sensitive" few lines of code.
3. yet another approach is to change the loss function to punish the high weights and encourage the model to use weights in a safe fp16 range, e.g. for t5
https://github.com/huggingface/transformers/pull/10956#issuecomment-820712267 - which may or may not work here and of course need to think how to add the extra component in a sensible way.

Then you can PR the changes and hopefully others will enjoy the fruit of your hard labour. Thank you!

Thank you for guiding me on how to debug the model and pointing out possible fixes. It took me some time to wrap my head around fp16. I think now I have a clear understanding on how to approach it. 

For now I made simple patches and will be running some more training and see how it goes. If I get some nice results, I'll post some summary here and do a PR.

A while back I also made a short study comparing bf16 and fp16, so it might be useful too to understand the limitations of fp16:
https://github.com/stas00/ml-ways/blob/master/numbers/bfloat16-vs-float16-study.ipynb