Fix qwen2vl float16 inference bug
Author: GeLee-Q
Pull request: yes
Labels: name: float16

# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

Fixes  [https://github.com/huggingface/transformers/issues/33294](https://github.com/huggingface/transformers/issues/33294)


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ ✅] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ✅] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?

### Infer res

This image shows a table comparing the performance of three precision modes—fp32, bf16, and fp16—in text generation:

fp32: Memory usage is 36629MiB, and it generated a description of a beach scene.
bf16: Memory usage is 44105MiB, and it generated a similar description.
fp16: An error occurred, and only exclamation marks were generated.
fp16 (improved version): Memory usage is 31609MiB, and it generated a more detailed description of the beach scene.
Each column displays the memory usage and the generated text or error information.

<img width="1563" alt="image" src="https://github.com/user-attachments/assets/264e0bea-0a17-492a-8e44-b2cb04ff7423">

### DeBug Process
Using the hook mechanism, I exported the input and output of each layer. When both the mean and sum are nan, I consider it to be abnormal.

```python
# fp16 hook 
Layer model.layers.0.self_attn.o_proj input - mean: nan, sum: nan, shape: torch.Size([1, 3602, 1536]), has inf: False, has nan: True, first 5 values: [-0.33251953125, 0.1829833984375, 0.064453125, -0.634765625, 0.5859375], last 5 values: [0.038909912109375, -0.08502197265625, 0.4794921875, 0.0762939453125, -0.1522216796875], 
Layer model.layers.0.self_attn.o_proj output - mean: nan, sum: nan, shape: torch.Size([1, 3602, 1536]), has inf: False, has nan: True, first 5 values: [0.0076446533203125, 0.019805908203125, 0.03533935546875, -0.045318603515625, -0.03057861328125], last 5 values: [-0.296875, 0.09539794921875, -0.0924072265625, 0.0084075927734375, -0.09539794921875], 

# fp32 hook 
Layer model.layers.0.self_attn.o_proj input - mean: 0.012129077687859535, sum: 67106.2109375, shape: torch.Size([1, 3602, 1536]), first 5 values: [-0.33251953125, 0.182861328125, 0.06439208984375, -0.63427734375, 0.5859375], last 5 values: [0.03431564196944237, -0.09479156881570816, 0.47655850648880005, 0.06808013468980789, -0.1560981571674347]
Layer model.layers.0.self_attn.o_proj output - mean: 0.008136402815580368, sum: 45016.046875, shape: torch.Size([3602, 1536]), first 5 values: [0.007521241903305054, 0.019816547632217407, 0.03532881289720535, -0.04532395675778389, -0.03064415045082569], last 5 values: [-0.2958013117313385, 0.09288226813077927, -0.08813470602035522, 0.007658728398382664, -0.09701263904571533]

```

- Further debugging in the forward process revealed that the attn_weights contained inf values. After applying softmax, nan values appeared, which subsequently caused all the following results to become nan.

```python
class Qwen2VLAttention(nn.Module):
		...
		attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
		...
```

###  Thoughts on bug:
- First, the problematic code is identical to that in LLaMA, so it’s unlikely that the issue is related to the way Torch is written.
- Second, LLaMA’s fp16 inference works without issues, making it difficult to suspect that the precision problem lies in the underlying operator called by torch.matmul.
- Lastly, I cautiously suspect that the open-source weights of Qwen-2VL may not be suitable for float16 inference, especially since the Qwen documentation indicates that[ fp16 is not recommended.](https://qwen.readthedocs.io/en/latest/inference/chat.html)

<img width="773" alt="image" src="https://github.com/user-attachments/assets/17265eeb-d3b1-48ae-a9dd-a7534d28067a">


### Training Comparison
**This modification is also effective for training.**
- code` llama-factory` https://github.com/hiyouga/LLaMA-Factory

#### fp32

```python
{"current_steps": 1, "total_steps": 100, "loss": 2.4317, "learning_rate": 1e-05, "epoch": 0.7272727272727273, "percentage": 1.0, "elapsed_time": "0:00:02", "remaining_time": "0:04:45"}
{"current_steps": 2, "total_steps": 100, "loss": 2.3595, "learning_rate": 2e-05, "epoch": 1.4545454545454546, "percentage": 2.0, "elapsed_time": "0:00:05", "remaining_time": "0:04:15"}
{"current_steps": 3, "total_steps": 100, "loss": 2.4036, "learning_rate": 3e-05, "epoch": 2.1818181818181817, "percentage": 3.0, "elapsed_time": "0:00:07", "remaining_time": "0:03:55"}
{"current_steps": 4, "total_steps": 100, "loss": 2.3881, "learning_rate": 4e-05, "epoch": 2.909090909090909, "percentage": 4.0, "elapsed_time": "0:00:09", "remaining_time": "0:03:42"}
{"current_steps": 5, "total_steps": 100, "loss": 2.3942, "learning_rate": 5e-05, "epoch": 3.6363636363636362, "percentage": 5.0, "elapsed_time": "0:00:11", "remaining_time": "0:03:34"}
```
![image](https://github.com/user-attachments/assets/f2523ce2-bf7c-474e-aaf7-06c2bf0002b2)

#### fp16

```python
{"current_steps": 1, "total_steps": 1000, "loss": 0.0, "learning_rate": 0.0, "epoch": 0.09195402298850575, "percentage": 0.1, "elapsed_time": "0:00:02", "remaining_time": "0:39:13"}
{"current_steps": 2, "total_steps": 1000, "loss": 0.0, "learning_rate": 0.0, "epoch": 0.1839080459770115, "percentage": 0.2, "elapsed_time": "0:00:04", "remaining_time": "0:35:30"}
{"current_steps": 3, "total_steps": 1000, "loss": 0.0, "learning_rate": 0.0, "epoch": 0.27586206896551724, "percentage": 0.3, "elapsed_time": "0:00:06", "remaining_time": "0:34:16"}
{"current_steps": 4, "total_steps": 1000, "loss": 0.0, "learning_rate": 0.0, "epoch": 0.367816091954023, "percentage": 0.4, "elapsed_time": "0:00:08", "remaining_time": "0:33:24"}
{"current_steps": 5, "total_steps": 1000, "loss": 0.0, "learning_rate": 0.0, "epoch": 0.45977011494252873, "percentage": 0.5, "elapsed_time": "0:00:09", "remaining_time": "0:32:48"}
```


![image](https://github.com/user-attachments/assets/2d8967e2-849a-45db-87ae-2fe1c7570d63)
#### fp16 this PR

```python
{"current_steps": 1, "total_steps": 1000, "loss": 4.5531, "learning_rate": 1.0000000000000002e-06, "epoch": 0.09195402298850575, "percentage": 0.1, "elapsed_time": "0:00:02", "remaining_time": "0:40:28"}
{"current_steps": 2, "total_steps": 1000, "loss": 4.5833, "learning_rate": 2.0000000000000003e-06, "epoch": 0.1839080459770115, "percentage": 0.2, "elapsed_time": "0:00:04", "remaining_time": "0:35:50"}
{"current_steps": 3, "total_steps": 1000, "loss": 4.1749, "learning_rate": 3e-06, "epoch": 0.27586206896551724, "percentage": 0.3, "elapsed_time": "0:00:06", "remaining_time": "0:34:13"}
{"current_steps": 4, "total_steps": 1000, "loss": 4.1741, "learning_rate": 4.000000000000001e-06, "epoch": 0.367816091954023, "percentage": 0.4, "elapsed_time": "0:00:08", "remaining_time": "0:33:19"}
{"current_steps": 5, "total_steps": 1000, "loss": 4.8369, "learning_rate": 5e-06, "epoch": 0.45977011494252873, "percentage": 0.5, "elapsed_time": "0:00:09", "remaining_time": "0:32:58"}
{"c
```

![image](https://github.com/user-attachments/assets/d5bcbaad-48c3-45a0-ae61-767a385647f7)


###  Personal Opinion on This Modification:
- This is just a trick that replaces all inf values resulting from torch.matmul with zero.
- This is not a fundamental solution to the problem, but it shows significant effects in fp16 inference and LoRA training. Feedback and appropriate suggestions from everyone are needed.

@ArthurZucker @zucchini-nlp @hiyouga @simonJJJ 

Comments:
> Hey! This indeed looks super useful, my 2 main concerns:
> 
> * speed: at inference in fp16 this will be slower
> 
> Otherwise this sounds very much reasonable. I think llama sometimes suffers from this as well

Are you suggesting that adding this operation will cause the inference speed to slow down in FP16? But if the precision is off, doesn't speed seem less important? I will test the full process inference speed tomorrow to compare. Additionally, regarding llama, I will test it using float16. This is a very interesting phenomenon.

> Approving, but let's make sure you run the slow tests actually! Broad support it probably more important than speed here!

Alright, I'll provide some data tomorrow.


> Approving, but let's make sure you run the slow tests actually! Broad support it probably more important than speed here!

@ArthurZucker Hello, my slow test has been completed. This operation incurs some overhead, but it's not significant. It accounts for 3.51% of the total CPU time and 0.6% of the total CUDA time across all operators. Trading this small overhead for improved accuracy is, in my opinion, acceptable overall.

### Comparison results
- input-pic  https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg
- input-text "Describe this image."

<img width="861" alt="image" src="https://github.com/user-attachments/assets/c6fab1cd-0ac8-4645-9419-c5880fe17f3b">

### profile
```
if query_states.dtype == torch.float16:
     attn_weights = torch.where(torch.isinf(attn_weights), torch.zeros_like(attn_weights), attn_weights)
```
This modification added three operations: aten::where, aten::isinf, and aten::zeros_like.
- aten::where:
    - Has high CPU usage, accounting for 0.98% of the total CPU time
    - Also has high CUDA usage, accounting for 0.55% of the total CUDA time
- aten::isinf:  Has the longest total CPU usage time (160.241ms), accounting for 1.70%, which is the highest among the three
- aten::zeros_like: Has relatively low CPU and CUDA usage

These three operations collectively consume approximately 3.51% of the total CPU time. From the nysys results, we can see that aten::where accounts for 0.6% of the total overhead among all operators.

#### torch profile 

Operation | Self CPU % | Self CPU | CPU total % | CPU total | CPU time avg | Self CUDA | Self CUDA % | CUDA total | CUDA time avg | CPU Mem | Self CPU Mem | CUDA Mem | Self CUDA Mem | # of Calls | Total GFLOPs
-- | -- | -- | -- | -- | -- | -- | -- | -- | -- | -- | -- | -- | -- | -- | --
aten::where | 0.45% | 42.024ms | 0.98% | 92.229ms | 27.416us | 24.981ms | 0.55% | 26.552ms | 7.893us | 0 b | 0 b | 8.39 Gb | 0 b | 3364 | --
aten::isinf | 0.14% | 13.013ms | 1.70% | 160.241ms | 49.335us | 0.000us | 0.00% | 33.928ms | 10.446us | 0 b | 0 b | 5.71 Gb | -7.46 Gb | 3248 | --
aten::zeros_like | 0.15% | 13.792ms | 0.83% | 78.403ms | 24.139us | 0.000us | 0.00% | 12.983ms | 3.997us | 0 b | 0 b | 8.38 Gb | -258.50 Kb | 3248 | --

#### nysys profile
Time (%) | Total Time (ns) | Num Calls | Avg (ns) | Med (ns) | Min (ns) | Max (ns) | StdDev (ns) | Name
-- | -- | -- | -- | -- | -- | -- | -- | --
0.6 | 26321648 | 3248 | 8104.0 | 2816.0 | 2175 | 624031 | 57133.6 | void at::native::vectorized_elementwise_kernel<(int)4, at::native::<unnamed>::where_kernel_impl(at:.

### test code 
```
...

def log_time_taken(start_time, step_name):
    end_time = time.time()
    logger.info(f"Time taken for {step_name}: {end_time - start_time:.4f} seconds")

...
logger.info("Starting model inference...")
inference_start = time.time()

output_ids = model.generate(
    **inputs,
    max_new_tokens=128,
    do_sample=False,
    top_k=None,
    top_p=None,
    temperature=None,
)

log_time_taken(inference_start, "Model Inference")
```

#### change dtype

```
# fp32, fp16, bf16 for torch_dtype
model = Qwen2VLForConditionalGeneration.from_pretrained(
    model_dir,
    torch_dtype=torch.bfloat16,
)
model.to(device)
```


> Hey! This indeed looks super useful, my 2 main concerns:
> 
> * speed: at inference in fp16 this will be slower
> 
> Otherwise this sounds very much reasonable. I think llama sometimes suffers from this as well

Can you provide an example of an issue with LLaMA? I've tested fp16 on my end and didn't encounter any particularly strange situations.

@ArthurZucker @zucchini-nlp @ydshieh 
Hello, I have completed the speed and accuracy tests. Is there anything else that needs to be done? I have already added the "run-slow" tag. Could you please trigger the action and then merge it in?

> Sorry for being late to re-review! Really nice numbers, thanks for deep diving. It is completely acceptable IMO! Even more so because it's on the `Qwen2VLAttention` (so not `Qwen2VLSdpaAttention` or Flash which would probably no suffer from this!)
> 
> Let's merge this!

@ArthurZucker  Thank you for the review. The CI has finished running. Could you please merge it?



Yeah sorry for the wait, wanted to compare a bit with #33317 which is quite similar.