OPT - Fix Softmax NaN in half precision mode
Author: younesbelkada
Pull request: yes

# What does this PR do?

Fix overflow / unstable operation issues when using large OPT models in half precision

- As it is done in [Megatron-DeepSpeed](https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/908dc9cb4b9717707241eaf8b92a986b2e251218/megatron/model/fused_softmax.py#L205), for large models it appears that you will have to first upcast the input to float32 before applying the Softmax function to avoid unexpected NaNs. This is because we use very large values (eg `-3.24e+38`) to mask the padded tokens. EDIT: it seems to we use correct values to mask padded tokens
- Linked issue: #17433 
- We'll probably need to re-compute the logits for slow tests but I am not sure

cc @patrickvonplaten @ArthurZucker @ydshieh @stas00 

Comments:
_The documentation is not available anymore as the PR was closed or merged._

Hi @younesbelkada 

```
expanded_attn_mask = _expand_mask(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1])
```
So when running in `half precision`, `_expand_mask` will use `torch.finfo(dtype)` with `dtype = inputs_embeds.dtype = fp16` and the `min` is `-65536`.

Am I missing anything here?

Is `fp32.min` used unexpectedly instead of `fp16.min` in this particular issue?

I have a PR #17306 for related issue. If using `-65536` has issue, then I need to hold on that PR to investigate.

Hi @ydshieh !
I think that you are right, when running in half precision I have `-65530` and not `-3.24e+38` in the attention mask as I said. But even with this mask I get NaNs on the padded hidden states for opt-1.3b, and upcasting the input to fp32 and casting back to fp16 seems to solve the issue for now 

> Hi @ydshieh ! I think that you are right, when running in half precision I have `-65530` and not `-3.24e+38` in the attention mask as I said. But even with this mask I get NaNs on the padded hidden states for opt-1.3b, and upcasting the input to fp32 and casting back to fp16 seems to solve the issue for now

Let me check - as if this is the case, the PR #17306 needs to find another way out 😢 

> I get NaNs on the padded hidden states for opt-1.3b, 

@younesbelkada 

- Could you point me which line in OPTModel you got `NaN  for padded hidden states`?
- Did you use the generation script in the linked issue, or you just run the model with some input ids? If it is the later case, could you provide the code snippet 🙏 please?

- I got NaNs exactly here: https://github.com/huggingface/transformers/blob/8f46ac98498dd47701971064617e00d7e723a98e/src/transformers/models/opt/modeling_opt.py#L217 - to fix it you can just do ` attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len).float() + attention_mask` [here](https://github.com/huggingface/transformers/blob/8f46ac98498dd47701971064617e00d7e723a98e/src/transformers/models/opt/modeling_opt.py#L214) and then `attn_weights = nn.functional.softmax(attn_weights, dim=-1).half()` [here](https://github.com/huggingface/transformers/blob/8f46ac98498dd47701971064617e00d7e723a98e/src/transformers/models/opt/modeling_opt.py#L217)
- Yes use the generation script provided in the issue, ie:
```
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, use_cache=True)
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.
greedy_output = model.generate(
    input_ids=batch["input_ids"].to("cuda"),
    attention_mask=batch["attention_mask"].to("cuda"),
    do_sample=False, top_k=0
)
```
Note also that everything works fine when `torch_dtype` is set to `torch.float32` or `torch.bfloat16` 

@younesbelkada 

The root cause is `-inf` is used here

https://github.com/huggingface/transformers/blob/8f46ac98498dd47701971064617e00d7e723a98e/src/transformers/models/opt/modeling_opt.py#L64

Change it to `mask = torch.full((tgt_len, tgt_len), torch.finfo(dtype).min)` **should** be fine.
(+/- inf * 0.0 will result NaN ).

### More details

With the above fix, there is still a minor issue. In 

https://github.com/huggingface/transformers/blob/8f46ac98498dd47701971064617e00d7e723a98e/src/transformers/models/opt/modeling_opt.py#L525

for batch index 0, we will see an `-inf`

```
tensor([[[[-65504.,    -inf, -65504., -65504., -65504., -65504., -65504.],
          [-65504., -65504., -65504., -65504., -65504., -65504., -65504.],
          [-65504., -65504.,      0., -65504., -65504., -65504., -65504.],
          [-65504., -65504.,      0.,      0., -65504., -65504., -65504.],
          [-65504., -65504.,      0.,      0.,      0., -65504., -65504.],
          [-65504., -65504.,      0.,      0.,      0.,      0., -65504.],
          [-65504., -65504.,      0.,      0.,      0.,      0.,      0.]]],
```

This is because we have `-65504` for causal mask + `-65504` due to (left) padding.
Regarding this part, we need to discuss with the team.

In general, we shouldn't have or use `-inf` (the only safe place to use it is immediately before the softmax).


Great! My suggestion is to mix both - we can force the attention mask to use -65504 for fp16 + upcast in fp32 and cast it back to fp16 after softmax for sanity check and avoid possible overflow issues. - Wdyt?

@stephenroller @suchenzang have you seen something similar in your training / inference runs? 

Also cc @patil-suraj - see issue. Would be nice to hear your opinion here

FYI, it can happen that during training you never use padding tokens. I may be mistaken but I know that for Bloom we do not train on padded batch inputs but on truncated sequences instead. 
Usually these issues can happen at inference time only!

Upcast to fp32 should never be required if masked tokens are masked with something that's not -inf. Upcast to fp32 is significant performance penalty. Single `-inf` value shouldn't be a problem as long as there are some non-zero values in the row, it would change output a little bit but that output is meaningless anyway, the whole row is masked out. 

Great thank you all for your comments and help! Following your advice I have added the changes proposed by @ydshieh - let me know if this works for you!

[This change](https://github.com/huggingface/transformers/blob/77162b94bddd51bb57c712e973e23eed2cd39971/src/transformers/models/opt/modeling_opt.py#L64) is also in #17306, but I am fine for a quick fix for `OPTModel`.

I would still like to point out that, although it is not useful for real usage of the model, leaving non-zero large negative values mixed with `-inf` to mask a whole sequence is not good for testing/debugging purpose -> but this could be addressed in another PR.



> [This change](https://github.com/huggingface/transformers/blob/77162b94bddd51bb57c712e973e23eed2cd39971/src/transformers/models/opt/modeling_opt.py#L64) is also in #17306, but I am fine for a quick fix for `OPTModel`.
> 
> I would still like to point out that, although it is not useful for real usage of the model, leaving non-zero large negative values mixed with `-inf` to mask a whole sequence is not good for testing/debugging purpose -> but this could be addressed in another PR.

Forgot to say, with current change, it's still possible to get `[-inf, -inf, dtype.min, dtype.min ...]` or `[-inf, -inf, -inf]` etc. after summing with the `attn_weights` (as mentioned, this depends the values in `attn_weights`). I will try to implement some processing in #17306 today.

We perform the upcast in our code, though we do it with softmax(dtype=torch.float32). It's very important.

That's an excellent point, Stephen! Thank you for that crucial reminder.

Indeed,  for pytorch ops that support accumulate dtype arg this approach makes things much more efficient than manual casting.

I remember I discovered that when optimizing the `LabelSmoother`:

https://github.com/huggingface/transformers/blob/7999ec125fc31428ed6879bf01bb013483daf704/src/transformers/trainer_pt_utils.py#L481

it made a huge difference.

> accumulate dtype arg this approach makes things much

For learning purpose, could you share why `using softmax(dtype=torch.float32)` is more efficient than explicit upcasting?

Because the op kernel does it automatically internally in a single operation by already accumulating in the correct dtype.

When you do it in 2 steps: `op(...).to(dtype=...)`, 2 additional memory copying operations have to happen to perform the casting.

@ngimel, did I explain that correctly? Thank you!

and it should be simple to benchmark the 2 cases to see the difference.

@Chillee, would `nvfuser` fuse explicit casting into the op's accumulate dtype automatically?


My original understanding of the process is like:

```
attn_scores =  attn_scores.to(torch.float32)
attn_prob =  nn.functional.softmax(attn_scores)
```

So I think the correct way should be:
```
attn_prob =  nn.functional.softmax(attn_scores, dtype=torch.float32)
```
right?

### Another question regarding dtype

After we get `attn_prob` in `float32`, should we cast it back to the target precision for the subsequential ops, like
```
attn_output = torch.bmm(attn_probs, value_states)
```
I am talking about the case where a user loads the models in fp16 and specify the inputs in fp16 too:
```
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
```
  - If we don't cast `attn_probs` back to the target type (here `fp16`)
    - it will fail (if `value_states` is `fp16`) for some op like `torch.bmm`
    - or will propagate the type fp32 for some simple ops (like `+`)

(I am not sure this is the correct/usual way to do inference in fp16, but this is what I see in the code snippet from the issue reporter)

I think the issue that this PR aims to address is not really about the upcast to float32. (@younesbelkada , right?)

It is mentioned in the PR description as a potential solution, but the original issue we want to address here comes from the fact that we get a sequence with all `-inf` as attention scores before `softmax`.

Maybe it it better to move the discussion(s) regarding the upcasting to another issue/PR page.

That's correct, the underlying issue is that for a row full of `-inf` softmax (by definition) produces `nan` (it's 0/0), and ideally that shouldn't be a problem because those fully masked row shouldn't participate in loss computation, but apparently they do and corrupt other values

> My original understanding of the process is like:
> 
> ```
> attn_scores =  attn_scores.to(torch.float32)
> attn_prob =  nn.functional.softmax(attn_scores)
> ```
> 
> So I think the correct way should be:
> 
> ```
> attn_prob =  nn.functional.softmax(attn_scores, dtype=torch.float32)
> ```
> 
> right?

It's correct. Just be aware that not all ops support this.

> ### Another question regarding dtype
> 
> After we get `attn_prob` in `float32`, should we cast it back to the target precision for the subsequential ops, like
> 
> ```
> attn_output = torch.bmm(attn_probs, value_states)
> ```
> 
> I am talking about the case where a user loads the models in fp16 and specify the inputs in fp16 too:
> 
> ```
> model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
> ```
> 
>     * If we don't cast `attn_probs` back to the target type (here `fp16`)
>       
>       * it will fail (if `value_states` is `fp16`) for some op like `torch.bmm`
>       * or will propagate the type fp32 for some simple ops (like `+`)
> 
> 
> (I am not sure this is the correct/usual way to do inference in fp16, but this is what I see in the code snippet from the issue reporter)

Yes, you definitely need to adjust the dtype to the one you expect.

In some cases it's enough to turn autocast off locally to have the whole ensemble automatically done in the right precision w/o any additional casting back and forth. For example see this workaround proposed for the t5 arch:

https://github.com/huggingface/transformers/pull/10956/files

```
    def forward(self, hidden_states):
        # many t5/mt5 models are trained in bfloat16 and don't do well under mixed precision (fp16).
        # It appears that it's enough to disable autocast for this FF layer to avoid inf/nan
        # problems for the whole model
        if torch.is_autocast_enabled():
            with torch.cuda.amp.autocast(enabled=False):
                return self._forward(hidden_states)
        else:
            return self._forward(hidden_states)
```


Hi @ydshieh, I am down for both solutions. We can either merge this PR as a quick patch for OPT or wait for #17306 to be merged! 
I can also open another PR to move the whole discussion around the upcasting issue (I think we need to address since it is done in the original OPT pipeline if I understood it right) - let me know what works best for you ;) ! 

@younesbelkada #17306 still needs some more reviews. So let's just see which PR is approved earlier and merge as it is.

@younesbelkada 

Since #17306 won't be merged at this moment, I guess you can try something like
(and see if the reviewers & pytorch experts like it 🙏 )

```
# change `-inf` to `dtype.min` to avoid `NaN` during `softmax`.
attn_scores = torch.max(attn_scores, torch.finfo(attn_scores.dtyte).min)

attn_prob = torch.nn.functional.softmax(attn_scores, ...) 
```

Hi all, 
I propose a fix in the latest commits as suggested by @ydshieh . To make it work I basically:
1- pre-process the att scores (suggestion by @ydshieh )
2- upcast the softmax in fp32 and cast it back to the original `dtype` (for consistency with what is done in the original implementation)
I also added a slow test to make sure these things do not happen in the future with OPT

I can also confirm that all slow test passes with this fix! Let me know what do you think ;) 

cc @ydshieh @patrickvonplaten 

@patil-suraj could you also take a quick look? :-)

Just a quick question before we merge. With this fix this issue here is solved: https://github.com/huggingface/transformers/issues/17433 ? What do the generations now give with this fix? Also I'm wondering a bit whether this is rather just because the model weights might be incorrect see: https://github.com/huggingface/transformers/issues/17653 

Should we maybe rather wait with this one until we have #17653 resolved? Or could we maybe run the examples of #17653 with this fix and see if we get better results?

@patrickvonplaten I can confirm that this implementation fixes #17433 - I added a test to make sure that this will never be produced again. The generations gave `Who are you? What do you want? What do you want? What do you` with this fix instead of `Who are you?  <\s>  <\s>  <\s>  <\s>` 
I think that we can merge this at least to fix this behavior + to make OPT implementation consistent with the one from Meta, since we do not upcast the softmax to float32 in our implementation. 
I think the problems with #17653 are related with the TP merging strategy 

Do we still need this fix @younesbelkada after having released the OPT fix ? Wondering because we do break the copy mechanism here which I'd like to avoid if possible - otherwise ok to merge I think

@patrickvonplaten I can confirm we are still observing this behaviour after the fix 
<img width="877" alt="Screenshot 2022-06-22 at 11 46 22" src="https://user-images.githubusercontent.com/49240599/174999015-074152ad-78b2-4c7c-a0b0-ac564115d022.png">
Could you share what do you mean by breaking the copy mechanism? Maybe we can avoid that by just retrieving the previous `dtype`without copying the old `attn_weights` as we do now


I can confirm the batched generations works fine now, can we merge?

Hi, as new code shows
```
if dtype_attn_weights == torch.float16:
    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(dtype_attn_weights)
else:
    attn_weights = nn.functional.softmax(attn_weights, dim=-1)
```
in modeling_opt.py at line 219 to 222, I wondered if it's equivalent to change these lines to
```
attn_weights = attn_weights-attn_weights.max(-1, keepdim=True).values
attn_weights = nn.functional.softmax(attn_weights, dim=-1)
```
In this case, we no longer need to upcast float16 to float32 and might speedup the training , and(or at leaest)  the inference?

cc @younesbelkada here

hey!
Sorry for responding here late, indeed it seems to be equivalent, ie the slow test that verifiies this specific issue successfully pass with your proposed change. 
Recall that the most crucial part to solve the issue is this line: `attn_weights = torch.max(attn_weights, torch.tensor(torch.finfo(attn_weights.dtype).min))` (at least as in today) and we decided to keep the `nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32)` to make the implementation consistent with original OPT as mentioned above.
However, your proposed changes break 2 `torch.fx` tests and I did not dig further into that, maybe if you open a PR we could discuss that into more details
Thanks!