diff --git a/vllm/model_executor/models/qwen3_omni_moe_thinker.py b/vllm/model_executor/models/qwen3_omni_moe_thinker.py index 460cbdcf66..52a5ccbe57 100755 --- a/vllm/model_executor/models/qwen3_omni_moe_thinker.py +++ b/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