From 1c917b275148fa6427ecf9eabb46bbd6df0f589e Mon Sep 17 00:00:00 2001 From: SYLAR <125541396+lishunyang12@users.noreply.github.com> Date: Sun, 23 Aug 2026 23:13:36 +0800 Subject: [PATCH] [Bugfix] Fix Qwen3-Omni AWQ quantization name mapping (#5687) Signed-off-by: lishunyang12 --- .../quantization/test_component_routing.py | 39 ++++++++++++++++++- .../models/qwen3_omni/quantization.py | 36 +++++++++++++++++ .../models/qwen3_omni/qwen3_omni.py | 36 ++++++++++++++++- .../models/qwen3_omni/qwen3_omni_code2wav.py | 5 ++- .../qwen3_omni/qwen3_omni_moe_talker.py | 4 ++ .../qwen3_omni/qwen3_omni_moe_thinker.py | 4 ++ vllm_omni/quantization/component_config.py | 11 ++++++ 7 files changed, 131 insertions(+), 4 deletions(-) create mode 100644 vllm_omni/model_executor/models/qwen3_omni/quantization.py diff --git a/tests/diffusion/quantization/test_component_routing.py b/tests/diffusion/quantization/test_component_routing.py index 9a306b6e4..943411779 100644 --- a/tests/diffusion/quantization/test_component_routing.py +++ b/tests/diffusion/quantization/test_component_routing.py @@ -9,6 +9,9 @@ import torch from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, ) +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( + CompressedTensorsConfig, +) from vllm.model_executor.models.utils import WeightsMapper from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import ( @@ -344,7 +347,41 @@ class TestTalkerVisualRouting: # =================================================================== -# 4. ComponentQuantizationConfig.resolve +# 4. ComponentQuantizationConfig.apply_vllm_mapper +# =================================================================== + + +class TestComponentApplyVllmMapper: + def test_maps_compressed_tensors_component_and_default(self): + thinker = CompressedTensorsConfig( + target_scheme_map={}, + ignore=["thinker.model.layers.0.mlp.gate"], + quant_format="pack-quantized", + ) + default = CompressedTensorsConfig( + target_scheme_map={}, + ignore=["talker.model.layers.0.mlp.gate"], + quant_format="pack-quantized", + ) + config = ComponentQuantizationConfig( + component_configs={"thinker": thinker, "code2wav": None}, + default_config=default, + ) + mapper = WeightsMapper( + orig_to_new_prefix={ + "thinker.model": "thinker.language_model.model", + "talker.model": "talker.language_model.model", + } + ).get_unstacked_mapper() + + config.apply_vllm_mapper(mapper) + + assert thinker.ignore == ["thinker.language_model.model.layers.0.mlp.gate"] + assert default.ignore == ["talker.language_model.model.layers.0.mlp.gate"] + + +# =================================================================== +# 5. ComponentQuantizationConfig.resolve # =================================================================== diff --git a/vllm_omni/model_executor/models/qwen3_omni/quantization.py b/vllm_omni/model_executor/models/qwen3_omni/quantization.py new file mode 100644 index 000000000..d870c7abd --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_omni/quantization.py @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Quantization mapping helpers for nested Qwen3-Omni stages.""" + +from vllm.model_executor.models.interfaces import SupportsQuant + +_OUTER_MAPPING_APPLIED = "_qwen3_omni_outer_mapping_applied" + + +def apply_outer_quant_config_mapping(model: SupportsQuant) -> None: + """Map checkpoint paths once at the unified model boundary.""" + SupportsQuant._maybe_apply_model_mapping(model) + if model.quant_config is not None: + setattr(model.quant_config, _OUTER_MAPPING_APPLIED, True) + + +def apply_nested_quant_config_mapping(model: SupportsQuant) -> None: + """Avoid applying a stage weight mapper twice to a mapped config.""" + quant_config = model.quant_config + if quant_config is None: + return + if not getattr(quant_config, _OUTER_MAPPING_APPLIED, False): + SupportsQuant._maybe_apply_model_mapping(model) + return + + # The outer model maps names, while the nested model still contributes + # fused-module metadata needed by quantization scheme matching. + if model.packed_modules_mapping is not None: + quant_config.packed_modules_mapping.update(model.packed_modules_mapping) + + +class Qwen3OmniNestedSupportsQuant(SupportsQuant): + """SupportsQuant variant for models initialized inside the stage wrapper.""" + + def _maybe_apply_model_mapping(self) -> None: + apply_nested_quant_config_mapping(self) diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py index ab524c578..71e2408d7 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -20,12 +20,22 @@ from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import ( from vllm.config import ModelConfig, VllmConfig from vllm.inputs import PromptType, TokensPrompt from vllm.logger import init_logger -from vllm.model_executor.models.interfaces import SupportsMRoPE, SupportsMultiModal, SupportsPP, SupportsRealtime +from vllm.model_executor.models.interfaces import ( + SupportsMRoPE, + SupportsMultiModal, + SupportsPP, + SupportsQuant, + SupportsRealtime, +) from vllm.model_executor.models.qwen3_asr_realtime import Qwen3ASRRealtimeBuffer from vllm.model_executor.models.qwen3_omni_moe_thinker import ( Qwen3OmniMoeConditionalGenerationMixin, ) -from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix +from vllm.model_executor.models.utils import ( + WeightsMapper, + init_vllm_registered_model, + maybe_prefix, +) from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import MultiModalFeatureSpec from vllm.sequence import IntermediateTensors @@ -39,6 +49,9 @@ from vllm_omni.data_entry_keys import Embeddings, HiddenStates, Ids, OmniPayload from vllm_omni.metrics import definitions as defs from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.model_executor.models.qwen3_omni.quantization import ( + apply_outer_quant_config_mapping, +) from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import ( Qwen3OmniMoeThinkerDummyInputsBuilder, Qwen3OmniMoeThinkerForConditionalGeneration, @@ -86,6 +99,7 @@ class Qwen3OmniMoeForConditionalGeneration( CustomProcessMixin, SupportsMRoPE, SupportsRealtime, + SupportsQuant, ): """ Unified Qwen3 Omni MoE model combining thinker, talker, and code2wav. @@ -99,6 +113,24 @@ class Qwen3OmniMoeForConditionalGeneration( Set `model_stage` in vllm_config to one of: "thinker", "talker", "code2wav" """ + # vLLM applies quantization-config name mapping before constructing this + # outer stage wrapper. Expose the final module paths here so checkpoint + # ignore lists (for example, the BF16 MoE routers in compressed-tensors + # checkpoints) match the nested thinker/talker modules at construction + # time. The keys intentionally omit a trailing dot so exact module names + # such as ``thinker.lm_head`` are mapped along with their parameters. + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + "thinker.lm_head": "thinker.language_model.lm_head", + "thinker.model": "thinker.language_model.model", + "talker.model": "talker.language_model.model", + } + ) + packed_modules_mapping = Qwen3OmniMoeThinkerForConditionalGeneration.packed_modules_mapping + + def _maybe_apply_model_mapping(self) -> None: + apply_outer_quant_config_mapping(self) + realtime_max_tokens = 64 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py index 3a85599d6..5c3129be9 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py @@ -28,11 +28,14 @@ from vllm.model_executor.models.utils import ( # type: ignore ) from vllm_omni.model_executor.models.common.snake_activation import SnakeBeta +from vllm_omni.model_executor.models.qwen3_omni.quantization import ( + Qwen3OmniNestedSupportsQuant, +) logger = init_logger(__name__) -class Qwen3OmniMoeCode2Wav(nn.Module): +class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant): """ Qwen3 Omni MoE Code2Wav - Converts num_quantizers-layer RVQ codec codes to audio waveform. diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py index 46401fb13..8337f8dad 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py @@ -18,6 +18,9 @@ from vllm.model_executor.models.utils import ( ) from vllm.sequence import IntermediateTensors +from vllm_omni.model_executor.models.qwen3_omni.quantization import ( + Qwen3OmniNestedSupportsQuant, +) from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_code_predictor_mtp import ( Qwen3OmniMoeTalkerCodePredictor, ) @@ -32,6 +35,7 @@ logger = init_logger(__name__) class Qwen3OmniMoeTalkerForConditionalGeneration( nn.Module, SupportsPP, + Qwen3OmniNestedSupportsQuant, ): """ Qwen3 Omni MoE Talker - Converts text to audio codec codes. 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 623c799d9..9b3198542 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 @@ -126,6 +126,9 @@ from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import ( _get_video_second_per_grid_t, _presampled_videos_hf_kwargs, ) +from vllm_omni.model_executor.models.qwen3_omni.quantization import ( + Qwen3OmniNestedSupportsQuant, +) from vllm_omni.quantization.component_config import ( PRE_QUANTIZED_METHODS, ComponentQuantizationConfig, @@ -1129,6 +1132,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( SupportsMRoPE, Qwen3OmniMoeConditionalGenerationMixin, SupportsTranscription, + Qwen3OmniNestedSupportsQuant, ): # PEFT stores the expert LoRA matrices with an explicit expert dimension. is_3d_moe_weight: bool = True diff --git a/vllm_omni/quantization/component_config.py b/vllm_omni/quantization/component_config.py index 638419689..4d74f7134 100644 --- a/vllm_omni/quantization/component_config.py +++ b/vllm_omni/quantization/component_config.py @@ -21,6 +21,9 @@ if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import ( QuantizeMethodBase, ) + from vllm.model_executor.models.utils import ( + WeightsMapper, + ) # Pre-quantized checkpoints (modelopt FP8/FP4/MXFP8/mixed) only quantize the @@ -112,6 +115,14 @@ class ComponentQuantizationConfig(QuantizationConfig): return self._components[comp_prefix] return self._default + def apply_vllm_mapper(self, hf_to_vllm_mapper: WeightsMapper) -> None: + """Apply a weight mapper to every routed quantization config.""" + for quant_config in self._components.values(): + if quant_config is not None: + quant_config.apply_vllm_mapper(hf_to_vllm_mapper) + if self._default is not None: + self._default.apply_vllm_mapper(hf_to_vllm_mapper) + def get_name(self) -> str: return "component"