From 7d9d8aee79d02d4616018ff277e9d32971c6da17 Mon Sep 17 00:00:00 2001
From: Siyuan Kong <101997501+ksiyuan@users.noreply.github.com>
Date: Thu, 21 May 2026 16:43:33 +0800
Subject: [PATCH] [Bugfix] Fix qwen2_5_omni weight loading (#3598)

Signed-off-by: ksiyuan <ksiyuan@umich.edu>
---
 .../qwen2_5_omni/test_weight_loading.py       | 49 +++++++++++++++++++
 .../models/qwen2_5_omni/qwen2_5_omni.py       |  1 -
 .../qwen2_5_omni/qwen2_5_omni_thinker.py      | 10 ++++
 3 files changed, 59 insertions(+), 1 deletion(-)
 create mode 100644 tests/model_executor/models/qwen2_5_omni/test_weight_loading.py

diff --git a/tests/model_executor/models/qwen2_5_omni/test_weight_loading.py b/tests/model_executor/models/qwen2_5_omni/test_weight_loading.py
new file mode 100644
index 000000000..1205dedb6
--- /dev/null
+++ b/tests/model_executor/models/qwen2_5_omni/test_weight_loading.py
@@ -0,0 +1,49 @@
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestQwen2_5OmniWeightLoading:
+    def test_qkv_weight_mapping(self):
+        mock_quant_description = {
+            "visual.blocks.0.attn.q.weight": {"quant_type": "FLOAT"},
+            "visual.blocks.0.attn.k.weight": {"quant_type": "FLOAT"},
+            "visual.blocks.0.attn.v.weight": {"quant_type": "FLOAT"},
+        }
+
+        prefix = "visual.blocks.0.attn.qkv"
+        packed_modules_mapping = {
+            "attn_akv_proj": [
+                "attn_q_proj",
+                "attn_k_proj",
+                "attn_v_proj",
+            ],
+            "qkv": [
+                "q",
+                "k",
+                "v",
+            ],
+        }
+
+        try:
+            proj_name = prefix.split(".")[-1]
+            if proj_name in packed_modules_mapping:
+                quant_type = None
+                shard_prefixes = [
+                    prefix.replace(proj_name, shard_proj_name) for shard_proj_name in packed_modules_mapping[proj_name]
+                ]
+                for shard_prefix in shard_prefixes:
+                    shard_quant_type = mock_quant_description[shard_prefix + ".weight"]
+
+                    if quant_type is None:
+                        quant_type = shard_quant_type
+                    elif shard_quant_type != quant_type:
+                        raise ValueError(
+                            f"Not all shards of {prefix} are quantized with same quant type."
+                            f"Shard {proj_name} uses {shard_quant_type}, but another shard"
+                            f"use {quant_type}. Please check quantization config."
+                        )
+            else:
+                quant_type = mock_quant_description[prefix + ".weight"]
+        except KeyError as e:
+            pytest.fail(f"KeyError was raised: {e}\n")
diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py
index 0b5d78560..712515f27 100644
--- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py
+++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py
@@ -81,7 +81,6 @@ class Qwen2_5OmniForConditionalGeneration(
             # Initialize thinker model (multimodal processing)
             self.thinker = init_vllm_registered_model(
                 vllm_config=vllm_config,
-                prefix=maybe_prefix(prefix, "thinker"),
                 hf_config=thinker_config,
                 # Use registry architecture key
                 architectures=["Qwen2_5OmniThinkerModel"],
diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py
index c4be88a88..931049d8d 100644
--- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py
+++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py
@@ -324,6 +324,16 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
             "gate_proj",
             "up_proj",
         ],
+        "attn_qkv_proj": [
+            "attn_q_proj",
+            "attn_k_proj",
+            "attn_v_proj",
+        ],
+        "qkv": [
+            "q",
+            "k",
+            "v",
+        ],
     }
 
     @classmethod