[Bugfix][Gemma4] Fix ModelOpt mixed-precision MoE config mapping (#48563)

Signed-off-by: wangqian <601731555@qq.com>
Co-authored-by: OpenAI Codex <codex@openai.com>

tests/quantization/test_modelopt.py

@@ -165,6 +165,38 @@ def test_modelopt_mixed_precision_does_not_quantize_unlisted_fused_sibling():
     assert config._resolve_quant_algo("model.layers.0.linear_attn.in_proj_ba") is None
 
 
+def test_modelopt_mixed_precision_composes_gemma4_mappers():
+    from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM
+    from vllm.model_executor.models.gemma4_mm import (
+        Gemma4ForConditionalGeneration,
+    )
+
+    config = _mixed_precision_config(
+        {
+            "model.language_model.layers.0.experts": {
+                "quant_algo": "NVFP4",
+                "group_size": 16,
+            },
+            "model.language_model.layers.1.moe.experts.gate_up_proj": {
+                "quant_algo": "NVFP4",
+                "group_size": 16,
+            },
+        }
+    )
+
+    config.apply_vllm_mapper(
+        Gemma4ForConditionalGeneration.hf_to_vllm_mapper.get_unstacked_mapper()
+    )
+    config.apply_vllm_mapper(Gemma4ForCausalLM.hf_to_vllm_mapper.get_unstacked_mapper())
+
+    expected_prefix = "language_model.model.layers.0.moe.experts"
+    assert set(config.quantized_layers) == {
+        expected_prefix,
+        "language_model.model.layers.1.moe.gate_up_proj",
+    }
+    assert config._resolve_quant_algo(expected_prefix) == "NVFP4"
+
+
 def test_modelopt_mixed_precision_infers_fused_gate_up_projection():
     from vllm.model_executor.layers.linear import LinearBase
 

vllm/model_executor/models/gemma4.py

@@ -84,6 +84,12 @@ from .utils import (
 
 logger = init_logger(__name__)
 
+_GEMMA4_EXPERT_PARENT_MAPPER = WeightsMapper(
+    orig_to_new_regex={
+        re.compile(r"(?<!\.moe)\.experts$"): ".moe.experts",
+    }
+)
+
 
 def _remap_gemma4_expert_weight_name(name: str) -> str:
     return re.sub(r"(?<!\.moe)\.experts\.(\d+)\.", r".moe.experts.\1.", name)
@@ -1508,7 +1514,7 @@ class Gemma4Model(nn.Module, EagleModelMixin):
 class Gemma4ForCausalLM(
     nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts, SupportsEagle3
 ):
-    hf_to_vllm_mapper = WeightsMapper(
+    hf_to_vllm_mapper = _GEMMA4_EXPERT_PARENT_MAPPER | WeightsMapper(
         orig_to_new_prefix={
             # Gemma4ForConditionalGeneration already loads the text stack
             # from `model.language_model.*`. We reuse that same checkpoint

vllm/model_executor/models/gemma4_mm.py

@@ -40,7 +40,10 @@ from vllm.inputs import MultiModalDataDict
 from vllm.logger import init_logger
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import ReplicatedLinear
-from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM
+from vllm.model_executor.models.gemma4 import (
+    _GEMMA4_EXPERT_PARENT_MAPPER,
+    Gemma4ForCausalLM,
+)
 from vllm.model_executor.models.module_mapping import MultiModelKeys
 from vllm.model_executor.models.transformers.utils import recursive_replace_linear
 from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -998,7 +1001,7 @@ class Gemma4ForConditionalGeneration(
     }
 
     # Maps checkpoint prefixes to vLLM module paths.
-    hf_to_vllm_mapper = WeightsMapper(
+    hf_to_vllm_mapper = _GEMMA4_EXPERT_PARENT_MAPPER | WeightsMapper(
         orig_to_new_prefix={
             # vision tower
             "model.vision_tower": "vision_tower",
@@ -1010,7 +1013,7 @@ class Gemma4ForConditionalGeneration(
             "model.language_model.": "language_model.model.",
             "lm_head.": "language_model.lm_head.",
             "model": "language_model.model",
-        }
+        },
     )
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):