From e24bddfb22f1c04ab28470043ac9a3585db95f64 Mon Sep 17 00:00:00 2001 From: leo <92263211+napleon-liu@users.noreply.github.com> Date: Mon, 20 Jul 2026 23:52:21 +0800 Subject: [PATCH] Fix Qwen3-Omni thinker fused MoE LoRA (#5191) Signed-off-by: leo --- .../qwen3_omni/test_qwen3_omni_moe_lora.py | 165 ++++++++++++++++++ .../qwen3_omni/qwen3_omni_moe_thinker.py | 12 ++ 2 files changed, 177 insertions(+) create mode 100644 tests/model_executor/models/qwen3_omni/test_qwen3_omni_moe_lora.py diff --git a/tests/model_executor/models/qwen3_omni/test_qwen3_omni_moe_lora.py b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_moe_lora.py new file mode 100644 index 000000000..e5fa38a6b --- /dev/null +++ b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_moe_lora.py @@ -0,0 +1,165 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression tests for Qwen3-Omni Thinker fused-MoE LoRA setup.""" + +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +from peft import LoraConfig as PeftLoraConfig +from peft import get_peft_model +from torch import nn +from transformers import PretrainedConfig +from vllm.config.lora import LoRAConfig +from vllm.lora.layers.fused_moe import FusedMoE3DWithLoRA +from vllm.lora.lora_model import LoRAModel +from vllm.lora.model_manager import LoRAModelManager +from vllm.lora.peft_helper import PEFTHelper + +from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import ( + Qwen3OmniMoeThinkerForConditionalGeneration, + _ensure_thinker_architecture, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_NUM_EXPERTS = 2 +_HIDDEN_SIZE = 4 +_INTERMEDIATE_SIZE = 4 +_LORA_RANK = 1 +_EXPERTS_MODULE = "language_model.model.layers.0.mlp.experts" +_PEFT_TARGET_PARAMETERS = [ + "thinker.model.layers.0.mlp.experts.gate_up_proj", + "thinker.model.layers.0.mlp.experts.down_proj", +] + + +class _TinyExperts(nn.Module): + """Expert parameters with the same 3-D layout as Qwen3-Omni.""" + + def __init__(self): + super().__init__() + self.gate_up_proj = nn.Parameter(torch.zeros(_NUM_EXPERTS, 2 * _INTERMEDIATE_SIZE, _HIDDEN_SIZE)) + self.down_proj = nn.Parameter(torch.zeros(_NUM_EXPERTS, _HIDDEN_SIZE, _INTERMEDIATE_SIZE)) + + +class _TinyQwen3OmniMoe(nn.Module): + """Minimal module tree needed for PEFT to save expert LoRA weights.""" + + def __init__(self, model_path: Path): + super().__init__() + self.thinker = nn.ModuleDict( + { + "model": nn.ModuleDict( + {"layers": nn.ModuleList([nn.ModuleDict({"mlp": nn.ModuleDict({"experts": _TinyExperts()})})])} + ) + } + ) + self.name_or_path = str(model_path) + self.config = {"model_type": "qwen3_omni_moe"} + + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + return inputs + + +class _CpuFusedMoE3DWithLoRA(FusedMoE3DWithLoRA): + """CPU harness for the real 3-D MoE LoRA weight methods.""" + + def __init__(self): + nn.Module.__init__(self) + self.moe_config = SimpleNamespace( + hidden_dim=_HIDDEN_SIZE, + num_local_experts=_NUM_EXPERTS, + num_experts=_NUM_EXPERTS, + intermediate_size_per_partition=_INTERMEDIATE_SIZE, + moe_parallel_config=SimpleNamespace( + tp_size=1, + tp_rank=0, + ep_rank=0, + use_ep=False, + ), + ) + self.device = torch.device("cpu") + self.tp_size = 1 + self.tp_rank = 0 + self._w13_slices = 1 + + +def _save_peft_adapter(tmp_path: Path) -> Path: + base_model_dir = tmp_path / "base_model" + base_model_dir.mkdir() + (base_model_dir / "config.json").write_text('{"model_type": "qwen3_omni_moe"}', encoding="utf-8") + + peft_model = get_peft_model( + _TinyQwen3OmniMoe(base_model_dir), + PeftLoraConfig( + r=_LORA_RANK, + lora_alpha=_LORA_RANK, + target_modules=[], + target_parameters=_PEFT_TARGET_PARAMETERS, + ), + ) + with torch.no_grad(): + for value, parameter in enumerate( + (parameter for parameter in peft_model.parameters() if parameter.requires_grad), + start=1, + ): + parameter.fill_(value) + + adapter_dir = tmp_path / "adapter" + peft_model.save_pretrained(adapter_dir) + return adapter_dir + + +def test_peft_expert_lora_loads_into_thinker(tmp_path: Path): + adapter_dir = _save_peft_adapter(tmp_path) + peft_helper = PEFTHelper.from_local_dir(str(adapter_dir), max_position_embeddings=128) + lora_model = LoRAModel.from_local_checkpoint( + str(adapter_dir), + {"experts"}, + peft_helper=peft_helper, + lora_model_id=1, + device="cpu", + dtype=torch.float32, + weights_mapper=Qwen3OmniMoeThinkerForConditionalGeneration.hf_to_vllm_mapper, + ) + + thinker_config = PretrainedConfig() + _ensure_thinker_architecture(thinker_config) + fused_moe = _CpuFusedMoE3DWithLoRA() + fused_moe.create_lora_weights( + max_loras=1, + lora_config=LoRAConfig( + max_lora_rank=_LORA_RANK, + max_loras=1, + lora_dtype=torch.float32, + ), + model_config=thinker_config, + ) + + manager = object.__new__(LoRAModelManager) + manager.is_pooling_model = False + manager._is_3d_moe_model = Qwen3OmniMoeThinkerForConditionalGeneration.is_3d_moe_weight + manager._stack_moe_lora_weights(lora_model, fused_moe, _EXPERTS_MODULE) + + expert_lora = lora_model.get_lora(_EXPERTS_MODULE) + assert expert_lora is not None + fused_moe.set_lora(0, expert_lora.lora_a, expert_lora.lora_b) + + assert fused_moe.adapter_enabled.tolist() == [1, 0] + for weights in ( + fused_moe.w13_lora_a_stacked[0], + fused_moe.w13_lora_b_stacked[0], + fused_moe.w2_lora_a_stacked[0], + fused_moe.w2_lora_b_stacked[0], + ): + assert torch.count_nonzero(weights[0]) == weights[0].numel() + + +def test_existing_thinker_architecture_is_preserved(): + config = PretrainedConfig(architectures=["CustomThinkerArchitecture"]) + + _ensure_thinker_architecture(config) + + assert config.architectures == ["CustomThinkerArchitecture"] diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py index fde493deb..65ec102a1 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py @@ -137,6 +137,14 @@ except (ImportError, ModuleNotFoundError): logger = init_logger(__name__) +_THINKER_ARCHITECTURE = "Qwen3OmniMoeThinkerForConditionalGeneration" + + +def _ensure_thinker_architecture(config: PretrainedConfig) -> None: + """Give the nested Thinker config the architecture required by MoE LoRA.""" + if not config.architectures: + config.architectures = [_THINKER_ARCHITECTURE] + class Qwen3Omni_VisionTransformer(_Qwen3Omni_VisionTransformer): """Subclass that fixes Qwen2_5_VisionAttention.forward() compatibility. @@ -1121,6 +1129,9 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( Qwen3OmniMoeConditionalGenerationMixin, SupportsTranscription, ): + # PEFT stores the expert LoRA matrices with an explicit expert dimension. + is_3d_moe_weight: bool = True + hf_to_vllm_mapper = WeightsMapper( orig_to_new_prefix={ "thinker.lm_head.": "language_model.lm_head.", @@ -1164,6 +1175,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( self.vllm_config = vllm_config # needed for torch compile forward context hf_config = vllm_config.model_config.hf_config thinker_config: Qwen3OmniMoeThinkerConfig = getattr(hf_config, "thinker_config", None) or hf_config + _ensure_thinker_architecture(thinker_config) quant_config = vllm_config.quant_config multimodal_config = vllm_config.model_config.multimodal_config self.config = thinker_config