[Bugfix] Fix qwen2_5_omni weight loading (#3598)

Signed-off-by: ksiyuan <ksiyuan@umich.edu>

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")

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"],

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