[LoRA] Add Qwen3-Omni multimodal LoRA support (#52786)
Signed-off-by: linitra24 <Joy25810@foxmail.com>
linitra24 committed
Aug 24, 2026 at 13:25 UTC
460c08bc8a525082f37b1ba4c8e70558e5aa8e9e
1 file changed
+40
-7
vllm/model_executor/models/qwen3_omni_moe_thinker.py
+40
-7
@@ -66,7 +66,11 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
66
from vllm.model_executor.models.module_mapping import MultiModelKeys
67
from vllm.model_executor.models.qwen2_audio import Qwen2AudioProcessingInfo
68
from vllm.multimodal import MULTIMODAL_REGISTRY
69
-from vllm.multimodal.inputs import MultiModalFeatureSpec, MultiModalKwargsItems
69
+from vllm.multimodal.inputs import (
70
+ MultiModalFeatureSpec,
71
+ MultiModalKwargsItem,
72
+ MultiModalKwargsItems,
73
+)
74
from vllm.multimodal.parse import AudioProcessorItems, MultiModalDataItems
75
from vllm.multimodal.processing.processor import (
76
MultiModalPromptUpdates,
@@ -84,6 +88,7 @@ from vllm.v1.attention.backends.registry import AttentionBackendEnum
88
from .interfaces import (
89
MultiModalEmbeddings,
90
SupportsEagle3,
91
+ SupportsLoRA,
92
SupportsMRoPE,
93
SupportsMultiModal,
94
SupportsPP,
@@ -202,13 +207,13 @@ class Qwen3OmniMoeAudioAttention(nn.Module):
207
208
self.scaling = self.head_dim**-0.5
209
205
- self.qkv = QKVParallelLinear(
210
+ self.qkv_proj = QKVParallelLinear(
211
hidden_size=self.embed_dim,
212
head_size=self.head_dim,
213
total_num_heads=self.num_heads,
214
total_num_kv_heads=self.num_heads,
215
bias=True,
211
- prefix=f"{prefix}.qkv",
216
+ prefix=f"{prefix}.qkv_proj",
217
)
218
219
self.out_proj = RowParallelLinear(
@@ -232,7 +237,7 @@ class Qwen3OmniMoeAudioAttention(nn.Module):
237
max_seqlen: torch.Tensor | None,
238
) -> torch.Tensor:
239
seq_length, _ = hidden_states.size()
235
- qkv, _ = self.qkv(hidden_states)
240
+ qkv, _ = self.qkv_proj(hidden_states)
241
q, k, v = qkv.chunk(3, dim=-1)
242
q = q.view(1, seq_length, -1, self.head_dim)
243
k = k.view(1, seq_length, -1, self.head_dim)
@@ -323,9 +328,9 @@ class Qwen3OmniMoeAudioEncoder(nn.Module):
328
329
hf_to_vllm_mapper = WeightsMapper(
330
orig_to_new_stacked={
326
- ".self_attn.q_proj.": (".self_attn.qkv.", "q"),
327
- ".self_attn.k_proj.": (".self_attn.qkv.", "k"),
328
- ".self_attn.v_proj.": (".self_attn.qkv.", "v"),
331
+ ".self_attn.q_proj.": (".self_attn.qkv_proj.", "q"),
332
+ ".self_attn.k_proj.": (".self_attn.qkv_proj.", "k"),
333
+ ".self_attn.v_proj.": (".self_attn.qkv_proj.", "v"),
334
}
335
)
336
@@ -1584,12 +1589,14 @@ class Qwen3OmniMoeConditionalGenerationMixin(Qwen2_5OmniConditionalGenerationMix
1589
class Qwen3OmniMoeThinkerForConditionalGeneration(
1590
nn.Module,
1591
SupportsMultiModal,
1592
+ SupportsLoRA,
1593
SupportsPP,
1594
SupportsMRoPE,
1595
SupportsEagle3,
1596
Qwen3OmniMoeConditionalGenerationMixin,
1597
SupportsTranscription,
1598
):
1599
+ is_3d_moe_weight: bool = True
1600
hf_to_vllm_mapper = WeightsMapper(
1601
orig_to_new_prefix={
1602
"thinker.lm_head.": "language_model.lm_head.",
@@ -1597,6 +1604,9 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
1604
"thinker.": "",
1605
"talker.": None,
1606
"code2wav.": None,
1607
+ # Adapter trained/saved directly from the thinker model
1608
+ "model.": "language_model.model.",
1609
+ "lm_head.": "language_model.lm_head.",
1610
}
1611
)
1612
@@ -1610,6 +1620,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
1620
"gate_proj",
1621
"up_proj",
1622
],
1623
+ "qkv": ["qkv"], # for visual encoder
1624
}
1625
1626
supported_languages = ISO639_1_SUPPORTED_LANGS
@@ -1631,6 +1642,8 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
1642
thinker_config: Qwen3OmniMoeThinkerConfig = (
1643
vllm_config.model_config.hf_config.thinker_config
1644
)
1645
+ # MoE LoRA uses the architecture to determine its weight layout.
1646
+ thinker_config.architectures = vllm_config.model_config.hf_config.architectures
1647
quant_config = vllm_config.quant_config
1648
multimodal_config = vllm_config.model_config.multimodal_config
1649
self.config = thinker_config
@@ -2299,3 +2312,23 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
2312
connector="visual.merger",
2313
tower_model=["visual.", "audio_tower."],
2314
)
2315
+
2316
+ def get_mm_lora_token_counts(
2317
+ self,
2318
+ *,
2319
+ modality: str,
2320
+ mm_kwargs: MultiModalKwargsItem | None,
2321
+ num_mm_embeds: int,
2322
+ ) -> tuple[int, int | None]:
2323
+ if modality in ("image", "video"):
2324
+ hf_config = self.config
2325
+ vision_config = hf_config.vision_config
2326
+ merge_size = vision_config.spatial_merge_size
2327
+ tower_tokens = num_mm_embeds * merge_size**2
2328
+ connector_tokens = num_mm_embeds
2329
+
2330
+ if modality == "audio":
2331
+ tower_tokens = num_mm_embeds
2332
+ connector_tokens = num_mm_embeds
2333
+
2334
+ return tower_tokens, connector_tokens