[LoRA] Add Qwen3-Omni multimodal LoRA support (#52786)

Signed-off-by: linitra24 <Joy25810@foxmail.com>

vllm/model_executor/models/qwen3_omni_moe_thinker.py

@@ -66,7 +66,11 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
 from vllm.model_executor.models.module_mapping import MultiModelKeys
 from vllm.model_executor.models.qwen2_audio import Qwen2AudioProcessingInfo
 from vllm.multimodal import MULTIMODAL_REGISTRY
-from vllm.multimodal.inputs import MultiModalFeatureSpec, MultiModalKwargsItems
+from vllm.multimodal.inputs import (
+    MultiModalFeatureSpec,
+    MultiModalKwargsItem,
+    MultiModalKwargsItems,
+)
 from vllm.multimodal.parse import AudioProcessorItems, MultiModalDataItems
 from vllm.multimodal.processing.processor import (
     MultiModalPromptUpdates,
@@ -84,6 +88,7 @@ from vllm.v1.attention.backends.registry import AttentionBackendEnum
 from .interfaces import (
     MultiModalEmbeddings,
     SupportsEagle3,
+    SupportsLoRA,
     SupportsMRoPE,
     SupportsMultiModal,
     SupportsPP,
@@ -202,13 +207,13 @@ class Qwen3OmniMoeAudioAttention(nn.Module):
 
         self.scaling = self.head_dim**-0.5
 
-        self.qkv = QKVParallelLinear(
+        self.qkv_proj = QKVParallelLinear(
             hidden_size=self.embed_dim,
             head_size=self.head_dim,
             total_num_heads=self.num_heads,
             total_num_kv_heads=self.num_heads,
             bias=True,
-            prefix=f"{prefix}.qkv",
+            prefix=f"{prefix}.qkv_proj",
         )
 
         self.out_proj = RowParallelLinear(
@@ -232,7 +237,7 @@ class Qwen3OmniMoeAudioAttention(nn.Module):
         max_seqlen: torch.Tensor | None,
     ) -> torch.Tensor:
         seq_length, _ = hidden_states.size()
-        qkv, _ = self.qkv(hidden_states)
+        qkv, _ = self.qkv_proj(hidden_states)
         q, k, v = qkv.chunk(3, dim=-1)
         q = q.view(1, seq_length, -1, self.head_dim)
         k = k.view(1, seq_length, -1, self.head_dim)
@@ -323,9 +328,9 @@ class Qwen3OmniMoeAudioEncoder(nn.Module):
 
     hf_to_vllm_mapper = WeightsMapper(
         orig_to_new_stacked={
-            ".self_attn.q_proj.": (".self_attn.qkv.", "q"),
-            ".self_attn.k_proj.": (".self_attn.qkv.", "k"),
-            ".self_attn.v_proj.": (".self_attn.qkv.", "v"),
+            ".self_attn.q_proj.": (".self_attn.qkv_proj.", "q"),
+            ".self_attn.k_proj.": (".self_attn.qkv_proj.", "k"),
+            ".self_attn.v_proj.": (".self_attn.qkv_proj.", "v"),
         }
     )
 
@@ -1584,12 +1589,14 @@ class Qwen3OmniMoeConditionalGenerationMixin(Qwen2_5OmniConditionalGenerationMix
 class Qwen3OmniMoeThinkerForConditionalGeneration(
     nn.Module,
     SupportsMultiModal,
+    SupportsLoRA,
     SupportsPP,
     SupportsMRoPE,
     SupportsEagle3,
     Qwen3OmniMoeConditionalGenerationMixin,
     SupportsTranscription,
 ):
+    is_3d_moe_weight: bool = True
     hf_to_vllm_mapper = WeightsMapper(
         orig_to_new_prefix={
             "thinker.lm_head.": "language_model.lm_head.",
@@ -1597,6 +1604,9 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
             "thinker.": "",
             "talker.": None,
             "code2wav.": None,
+            # Adapter trained/saved directly from the thinker model
+            "model.": "language_model.model.",
+            "lm_head.": "language_model.lm_head.",
         }
     )
 
@@ -1610,6 +1620,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
             "gate_proj",
             "up_proj",
         ],
+        "qkv": ["qkv"],  # for visual encoder
     }
 
     supported_languages = ISO639_1_SUPPORTED_LANGS
@@ -1631,6 +1642,8 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
         thinker_config: Qwen3OmniMoeThinkerConfig = (
             vllm_config.model_config.hf_config.thinker_config
         )
+        # MoE LoRA uses the architecture to determine its weight layout.
+        thinker_config.architectures = vllm_config.model_config.hf_config.architectures
         quant_config = vllm_config.quant_config
         multimodal_config = vllm_config.model_config.multimodal_config
         self.config = thinker_config
@@ -2299,3 +2312,23 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
             connector="visual.merger",
             tower_model=["visual.", "audio_tower."],
         )
+
+    def get_mm_lora_token_counts(
+        self,
+        *,
+        modality: str,
+        mm_kwargs: MultiModalKwargsItem | None,
+        num_mm_embeds: int,
+    ) -> tuple[int, int | None]:
+        if modality in ("image", "video"):
+            hf_config = self.config
+            vision_config = hf_config.vision_config
+            merge_size = vision_config.spatial_merge_size
+            tower_tokens = num_mm_embeds * merge_size**2
+            connector_tokens = num_mm_embeds
+
+        if modality == "audio":
+            tower_tokens = num_mm_embeds
+            connector_tokens = num_mm_embeds
+
+        return tower_tokens, connector_tokens