OPT produce NaN during batched generation
Author: shijie-wu
Pull request: no
Labels: name: bug

### System Info

* transformers==4.19.2
* PyTorch (GPU?): 1.11.0+cu102 (True)
* GPUs: single V100


### Who can help?

@LysandreJik, @younesbelkada

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

```py
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# I have tested and the error happens to opt-1.3b, opt-2.7b, opt-6.7b, and opt-13b.
# opt-125m and opt-350m seems to work fine.
# I haven't tested opt-30b.
model_name = "facebook/opt-1.3b"
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
tokenizer.padding_side = "left"
# It works when torch_dtype=torch.float32
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
model = model.eval().to("cuda")

batch = tokenizer(
    ["Who are you?", "Joe Biden is the president of"],
    padding=True, return_tensors="pt"
)

# It produces NaN in the early layers for the first sequence.
# I check the pattern, and NaN first appears in the padded token position.
model.generate(
    input_ids=batch["input_ids"].to("cuda"),
    attention_mask=batch["attention_mask"].to("cuda"),
    do_sample=True, max_new_tokens=32
) 
```

### Expected behavior

The generation under fp16 should be close to fp32.


Comments:
Hi @shijie-wu ! Thanks for pointing out the issue,
It appears that in our current implementation we use a naive Softmax function that is applied to the attention scores. When passing the attention scores combined with the attention mask we use very large values to mask padding tokens such as `torch.inf` or `-3.24e+38` on the softmax function. (EDIT: see #17437 - we use correct padding values and not `-3.24e+38`)
It seems that this sometimes leads to unstable operations and results in having NaNs when using half precision mode, but only for large models. 
I think that the correct workaroud is to upcast the attention scores to float32 before summing it with the attention mask, apply the Softmax then cast it back to the input dtype as it is done in [Megatron-DeepSpeed](https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/908dc9cb4b9717707241eaf8b92a986b2e251218/megatron/model/fused_softmax.py#L204) for eg. 
A proper fix will be addressed in #17437 but a quick and dirty solution would be to use `bfloat16` instead of `float16`. At least it worked with `opt-1.3b` but I don't know if it will work with larger models.
Let me know if this helps!

Hi @younesbelkada ! Thank you for the quick response! I will follow https://github.com/huggingface/transformers/pull/17437. A fix for fp16 would be great as only A100 class support bf16 AFAIK.

Out of curiosity, do you have any intuition on why it only impacts larger models but not smaller models? From https://github.com/huggingface/transformers/pull/17437, it seems to me that it would impact smaller models as well?

Thanks for the comment! 
We will try to have a patch to fix that for fp16 asap I guess ;) Curious to know if the proposed PR will fix your issue (you can checkout the PR and build it from source if you have time)!

Regarding your second question - I totally agree with you - it should also not work on small models. It is just an intuition but possibly the number of heads and/or hidden dimension are impacting that (since it is the only thing that differs between `opt-125m` and `opt-1.3b` in the first layers). I would wait for the team's comments to see if they have better intuition on that! 

Interesting! Thanks a lot for reporting this @shijie-wu 

@shephenroller @suchenzang do you have any insight here maybe?

Also cc @patil-suraj, I remember we had a similar problem with GPT-Neo/GPT-J no? Was the solution to force the last computation of the logits to be in fp32?

Hi @shijie-wu !
After discussing with @ydshieh , it appears that this is completely independent from the model size but it just happened by luck that the logits before the softmax were negative only for large models - therefore causing an overflow that leaded to NaNs. 
I am quoting his answer here:

> And back to your question about "why only large model":
This is about the weights and inputs. In the generation script provided in the issue, when I run it, in second time passing Attention Layer, there is some point we get [-16.x, -16.x, ....] as attn_weights , and attn_mask  as [-65504, -65504, -inf, -inf, ...]
The 2 -65504 from left padding, -inf  from causal mask .
However in fp16, -65504 + -16 = -inf . So we get a batch index with all -inf  as input to softmax , and get outputs NaN`

Very interesting!

FYI (if you want to know more details)

https://github.com/huggingface/transformers/pull/17306#issuecomment-1138660341

> Also cc @patil-suraj, I remember we had a similar problem with GPT-Neo/GPT-J no? Was the solution to force the last computation of the logits to be in fp32?

@patrickvonplaten I dunno about GPT-Neo/GPT-J but  can vouch for the fact that the same thing is happening with GPT-NeoX currently.

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.

The PR https://github.com/huggingface/transformers/pull/17437 has been merged, @shijie-wu could you confirm this fixed your issue?

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.