attention weight clipping
Author: StevenSong
Pull request: no

### Feature request

If the attention weight overflows, probably when using float16 during mixed precision training,  clip the weight to some configurable value.

### Motivation

I’m training a `gpt2` model with auto mixed precision with `torch.amp.autocast` and I noticed I’m running into `nan` loss values during training. I tracked the source of the `nan` to a `softmax` computation where there’s a single `inf` in the input to the `softmax`. The `inf` is coming from a matrix multiply of the query and key matrices to calculate attention weights. this line: https://github.com/huggingface/transformers/blob/4b26a61631b8fd30f845cf08ebcc5ed65fe83c9b/src/transformers/models/gpt2/modeling_gpt2.py#L184.
Specifically the dot product of two vectors from query/key overflows the `float16` dtype.

### Your contribution

Would a simple `torch.clamp` call work/be correct?

Comments:
Hi!

There are a few places in the library having something below

```python
        # clamp inf values to enable fp16 training
        if hidden_states.dtype == torch.float16:
            max_dtype = torch.finfo(hidden_states.dtype).max
            clamp_value = torch.where(torch.isinf(hidden_states).any(), max_dtype - 1000, max_dtype)
            hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
```

You can try similar thing locally :-)

@ydshieh thank you for the example, that's effectively exactly what I'm looking for. 

My only concern is that implementing this clamping locally is far more overhead for an end user than implementing it in the codebase. My reasoning for this is because this inf is not really returned outside of the model forward call as it is passed into a softmax in the GPT2Attention block, and ultimately the user only sees the nans propagated by the inf in the softmax. So in order to enable this clamping, a user would have to override the forward call of `GPT2Attention` and ultimately subclass `GPT2Attention`, `GPT2Block`, `GPT2Model`, and whichever class they're using which contains the base model (eg `GPT2LMHeadModel`).

I understand @StevenSong . But it's better for you to try locally first to see if it actually solves the issue, or something more have to be done (like same fix at different places).

In theory, this kind of `overflows` can happen everywhere. In practice, we proably just need to add clamp at one or two places.
Also I would need to discuss internally to see if we want to do this change for a long-existing model like `gpt2`.

Just for my debugging, I ended up just modifying the source file at `models/gpt2/modeling_gpt2.py` and inserting the below chunk between these lines: https://github.com/huggingface/transformers/blob/4b26a61631b8fd30f845cf08ebcc5ed65fe83c9b/src/transformers/models/gpt2/modeling_gpt2.py#L203-L205

`inf` clamping chunk, same as what was suggested above:
```python
        if attn_weights.dtype == torch.float16:
            max_dtype = torch.finfo(attn_weights.dtype).max
            clamp_value = torch.where(torch.isinf(attn_weights).any(), max_dtype - 1000, max_dtype)
            attn_weights = torch.clamp(attn_weights, min=-clamp_value, max=clamp_value)
```

and on my test case, I can confirm that this results in non-nan loss and non-nan logits. I can continue with my training loop with no errors from backprop and the next batch also returns successfully

Thanks!

Oh, @StevenSong 

It's for attn weight. I just remembered that for it, it's starndard to cast it to fp32 instead of staying in fp16 and use claming.
Could you try if the following works in your case 🙏 ? Thanks a lot!

```python
        # upcast to fp32 if the weights are in fp16. Please see https://github.com/huggingface/transformers/pull/17437
        if attn_weights.dtype == torch.float16:
            attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(torch.float16)
        else:
            attn_weights = nn.functional.softmax(attn_weights, dim=-1)
```





Hi @ydshieh,

I did notice this upcasting was already being done implicitly when adding the `attention_mask` to `attn_weights` at this line:

https://github.com/huggingface/transformers/blob/4b26a61631b8fd30f845cf08ebcc5ed65fe83c9b/src/transformers/models/gpt2/modeling_gpt2.py#L207

As mentioned in the referenced PR (#17437), `attention_mask` is filled with very large negative values. edit: I guess this also depends on `attention_mask` being passed

But this is already far after the `inf` is introduced at the matmul of the query and key matrices together. I guess if those are cast to fp32, that would probably also fix it?

Yes, the large negative value depends on the dtype


https://github.com/huggingface/transformers/blob/4b26a61631b8fd30f845cf08ebcc5ed65fe83c9b/src/transformers/models/gpt2/modeling_gpt2.py#L822-L823.

If we upcast to fp32 as mentioned in my previous comment, it should be fine. This is done in a few models, like `opt` or `xglm`. 

See also #17437

> Yes, the large negative value depends on the dtype
> 
> https://github.com/huggingface/transformers/blob/4b26a61631b8fd30f845cf08ebcc5ed65fe83c9b/src/transformers/models/gpt2/modeling_gpt2.py#L822-L823
> 
> .


Interesting, I'd have expected `attention_mask` to be `-65k` but I'm seeing `-3.4028e+38` as in fp32, yet my `attn_weights` are in fp16. Is this because of `autocast` only casting some but not all tensors ie query/key are cast to fp16 but the model dtype is still fp32?

> If we upcast to fp32 as mentioned in my previous comment, it should be fine. This is done in a few models, like `opt` or `xglm`.
> 
> See also #17437

I agree, upcasting to fp32 should resolve the issue but I think it needs to be done earlier, at the level of query/key matrices. Otherwise the `inf` would just be upcast to fp32, no?

> Is this because of autocast only casting some but not all tensors ie query/key are cast to fp16 but the model dtype is still fp32?

I didn't use this before personally, but from torch doc

> where some operations use the torch.float32 (float) datatype and other operations use lower precision floating point datatype

it looks the same as you mentioned.

Regarding:

>  it needs to be done earlier, at the level of query/key matrices. Otherwise the inf would just be upcast to fp32, no?

The `inf` you see on query/key matrices might be also a consequence of the computation on attn_weight or its softmax in an earlier step. It's not easy to say for sure where the problem starts to accumulate (even they might not cause failure at that early time).

Let's try the more standard approaces where people suggest to using fp32 for softmax (and let's upcast before this while adding attention mask) and see how things go 🤗 .


Got to say my above comment is based on my experience on FP16, not with autocast. From your description, you mentioned the `attn_weights + attention_mask` would already be in FP32. It's good idea to double check this (what's the dtype of this step's output) and if the following softmax takes places in FP32 too. 

However, it doesn't hurt to give it a try 🤗 

apologies for the late reply to this thread, here's what I've tried and what's worked/not worked (and by works, I mean if `inf` and `nan` no longer appear in `attn_weights` and loss is non-`nan`):

1. upcasting at the softmax call (see below for code): this does NOT work as `attn_weights.dtype` is no longer fp16 after adding `attention_mask` in fp32. so the `inf` is still passed to softmax and we get `nan`s.
```python
    def _attn(self, query, key, value, attention_mask=None, head_mask=None):
        [...]
        if attention_mask is not None:
            # Apply the attention mask
            attn_weights = attn_weights + attention_mask

        if attn_weights.dtype == torch.float16:
            attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(torch.float16)
        else:
            attn_weights = nn.functional.softmax(attn_weights, dim=-1)
```
2. explicitly upcasting before summing: this does NOT work as `attn_weights` already contains the `inf` and so is simply upcasting `inf` from fp16 to fp32.
```python
    def _attn(self, query, key, value, attention_mask=None, head_mask=None):
        [...]
        attn_weights = attn_weights.to(torch.float32)
        if attention_mask is not None:
            # Apply the attention mask
            attn_weights = attn_weights + attention_mask

        attn_weights = nn.functional.softmax(attn_weights, dim=-1)
```
3. explicitly upcasting query and key matrices: this also does NOT work. this is because even though query and key are indeed in fp32, within the `autocast` context, `torch.matmul` downcasts back to fp16! 
```python
    def _attn(self, query, key, value, attention_mask=None, head_mask=None):
        query = query.to(torch.float32)
        key = key.to(torch.float32)
        attn_weights = torch.matmul(query, key.transpose(-1, -2))
        [...]
```
4. explicitly upcasting query and key with enforced fp32: the reason I'm so insistent on upcasting query and key is because I've already found that the line which produces the `inf` is the matmul between the query and key matrices, as I mentioned in the original post. However, as seen in the above attempt, the `autocast` context does not respect the explicit cast. So we need to disable the autocast context, if it exists ([see relevant docs](https://pytorch.org/docs/stable/notes/amp_examples.html#autocast-and-custom-autograd-functions)). Thus `attn_weights` is finally in fp32 as the product of two fp32 matrices.
```python
    def _attn(self, query, key, value, attention_mask=None, head_mask=None):
        query = query.to(torch.float32)
        key = key.to(torch.float32)
        if torch.is_autocast_enabled():
            with torch.amp.autocast(device_type=query.device.type, enabled=False):
                attn_weights = torch.matmul(query, key.transpose(-1, -2))
        else:
            attn_weights = torch.matmul(query, key.transpose(-1, -2))
```

I'd also put forth that upcasting the query and key vectors to fp32 is a generalizable solution, as `attn_weights` is then always fp32 and all subsequent operations with `attn_weights` in the attention block are also implicitly upcast to fp32. It can always be downcast later, in fact it seems like this was already considered in this same function at this line:
https://github.com/huggingface/transformers/blob/fe861e578f50dc9c06de33cd361d2f625017e624/src/transformers/models/gpt2/modeling_gpt2.py#L211-L213


Is there a good way for me to share my example I've been debugging? I've pared it down to a single script with a specific batch which results in the `inf`/`nan`, the base model is on the hub already, and I'm specifically doing prompt tuning so there's only something like extra ~20K params for the example to work

@StevenSong 

- thank you for the hard work on doing experiments! ❤️ 
- You can post the script as in a comment, or maybe create a colab notebook and share with us 🙏 

Also, I haven't asked (I think) previously: could you share us the full error log (I know it's inf thing, but would like to see the log). Thank you 🙏 




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.