Qwen2vl float16 inference bug in naive attention
Author: kjohew
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.46.2
- Platform: Linux-5.15.0-125-generic-x86_64-with-glibc2.35
- Python version: 3.10.15
- Huggingface_hub version: 0.26.3
- Safetensors version: 0.4.5
- Accelerate version: 1.1.1
- Accelerate config:    not found
- PyTorch version (GPU?): 2.5.1+cu124 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: <fill in>
- Using GPU in script?: <fill in>
- GPU type: NVIDIA A10

### Who can help?

@ArthurZucker @GeLee-Q 

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Hi @ArthurZucker @GeLee-Q , thanks for #33312, which has resolved the inference bug for qwen2vl in float16. However, it still has some flaws. Although the current code solves the numerical overflow issue, it may **break causality of the autoregressive model**. See details [here](https://github.com/huggingface/transformers/issues/35151#issuecomment-2537723488).

### Expected behavior

Current: [Code](https://github.com/huggingface/transformers/blob/c8c8dffbe45ebef0a8dba4a51024e5e5e498596b/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py#L582)
```
        if attention_mask is not None:  # no matter the length, we just slice it
            causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
            attn_weights = attn_weights + causal_mask

        # Fix precision issues in Qwen2-VL float16 inference
        # Replace inf values with zeros in attention weights to prevent NaN propagation
        if query_states.dtype == torch.float16:
            attn_weights = torch.where(torch.isinf(attn_weights), torch.zeros_like(attn_weights), attn_weights)
```


Fixed: Replace ±inf with zero before adding attn_weights with causal_mask. (prevent converting -inf, which for causality, to zero.)
```
        # Fix precision issues in Qwen2-VL float16 inference
        # Replace inf values with zeros in attention weights to prevent NaN propagation
        if query_states.dtype == torch.float16:
            attn_weights = torch.where(torch.isinf(attn_weights), torch.zeros_like(attn_weights), attn_weights)

        if attention_mask is not None:  # no matter the length, we just slice it
            causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
            attn_weights = attn_weights + causal_mask
```

Comments:
The previous modification might not have covered all scenarios. Could you provide an example to illustrate the shortcomings? If this modification is fine, please feel free to submit a PR.

Hi @GeLee-Q , there are two places that can easily generate ±inf:
1. `torch.matmul(query_states, key_states.transpose(2, 3))`
2. `attn_weights + causal_mask`
But only these inf generated in the first case is the one we want to set to zero. In the second case, A negative attn_weight added to -65504 (causal masked position) results in -inf, which we want to keep as is.

Here is a example:
```python
attn_weights = torch.tensor(
    [
        [-1, 100, float("inf")],
        [-1, -1, -100],
        [-1, -1, -1],
    ],
    dtype=torch.float16
)

causal_mask = torch.tensor(
    [
        [0., -65504., -65504.],
        [0.,      0., -65504.],
        [0.,      0.,      0.],
    ],
    dtype=torch.float16
)
```
According to the current code, we will obtain:
```python
attn_weights = attn_weights + causal_mask
attn_weights = torch.where(torch.isinf(attn_weights), torch.zeros_like(attn_weights), attn_weights)
print(attn_weights)
# tensor([[-1.0000e+00, -6.5408e+04,  0.0000e+00],
#         [-1.0000e+00, -1.0000e+00,  0.0000e+00],
#         [-1.0000e+00, -1.0000e+00, -1.0000e+00]], dtype=torch.float16)

attn_weights = torch.softmax(attn_weights, dim=-1)
print(attn_weights)
# tensor([[0.2690, 0.0000, 0.7310],
#         [0.2119, 0.2119, 0.5762],
#         [0.3333, 0.3333, 0.3333]], dtype=torch.float16)
```
We can see 0.7310 in attn_weights[0, 2] and 0.5762 in attn_weights[1, 2], which violates the causality of the autoregressive model.

In the fixed version, we not only solve the overflow problem introduced by torch.matmul, but also preserve the causality of llm.
```python
attn_weights = torch.where(torch.isinf(attn_weights), torch.zeros_like(attn_weights), attn_weights)
attn_weights = attn_weights + causal_mask
print(attn_weights)
# tensor([[-1.0000e+00, -6.5408e+04, -6.5504e+04],
#         [-1.0000e+00, -1.0000e+00,        -inf],
#         [-1.0000e+00, -1.0000e+00, -1.0000e+00]], dtype=torch.float16)

attn_weights = torch.softmax(attn_weights, dim=-1)
print(attn_weights)
# tensor([[1.0000, 0.0000, 0.0000],
#         [0.5000, 0.5000, 0.0000],
#         [0.3333, 0.3333, 0.3333]], dtype=torch.float16)
```






I agree with your fix - by adjusting the order of inf handling and mask addition, it effectively addresses both the numerical overflow issue and maintains model causality. As the contributor who solved this issue, I hope you'll consider submitting a PR when you have time. Thanks for sharing!

Hi @GeLee-Q, thanks for your response. I'm also very pleased that we have reached an agreement on the solution to fix this issue. I will submit a PR as soon as possible to address the issue and include comments in the code, referencing both these two related issues to facilitate future understanding and tracking.

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.