[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>
wangqian committed
Jul 20, 2026 at 21:39 UTC
0a5069e4e30996866d9d3cc5551f784b85408bb0
3 files changed
+45
-4
tests/quantization/test_modelopt.py
+32
@@ -165,6 +165,38 @@ def test_modelopt_mixed_precision_does_not_quantize_unlisted_fused_sibling():
165
assert config._resolve_quant_algo("model.layers.0.linear_attn.in_proj_ba") is None
166
167
168
+def test_modelopt_mixed_precision_composes_gemma4_mappers():
169
+ from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM
170
+ from vllm.model_executor.models.gemma4_mm import (
171
+ Gemma4ForConditionalGeneration,
172
+ )
173
+
174
+ config = _mixed_precision_config(
175
+ {
176
+ "model.language_model.layers.0.experts": {
177
+ "quant_algo": "NVFP4",
178
+ "group_size": 16,
179
+ },
180
+ "model.language_model.layers.1.moe.experts.gate_up_proj": {
181
+ "quant_algo": "NVFP4",
182
+ "group_size": 16,
183
+ },
184
+ }
185
+ )
186
+
187
+ config.apply_vllm_mapper(
188
+ Gemma4ForConditionalGeneration.hf_to_vllm_mapper.get_unstacked_mapper()
189
+ )
190
+ config.apply_vllm_mapper(Gemma4ForCausalLM.hf_to_vllm_mapper.get_unstacked_mapper())
191
+
192
+ expected_prefix = "language_model.model.layers.0.moe.experts"
193
+ assert set(config.quantized_layers) == {
194
+ expected_prefix,
195
+ "language_model.model.layers.1.moe.gate_up_proj",
196
+ }
197
+ assert config._resolve_quant_algo(expected_prefix) == "NVFP4"
198
+
199
+
200
def test_modelopt_mixed_precision_infers_fused_gate_up_projection():
201
from vllm.model_executor.layers.linear import LinearBase
202
vllm/model_executor/models/gemma4.py
+7
-1
@@ -84,6 +84,12 @@ from .utils import (
84
85
logger = init_logger(__name__)
86
87
+_GEMMA4_EXPERT_PARENT_MAPPER = WeightsMapper(
88
+ orig_to_new_regex={
89
+ re.compile(r"(?<!\.moe)\.experts$"): ".moe.experts",
90
+ }
91
+)
92
+
93
94
def _remap_gemma4_expert_weight_name(name: str) -> str:
95
return re.sub(r"(?<!\.moe)\.experts\.(\d+)\.", r".moe.experts.\1.", name)
@@ -1508,7 +1514,7 @@ class Gemma4Model(nn.Module, EagleModelMixin):
1514
class Gemma4ForCausalLM(
1515
nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts, SupportsEagle3
1516
):
1511
- hf_to_vllm_mapper = WeightsMapper(
1517
+ hf_to_vllm_mapper = _GEMMA4_EXPERT_PARENT_MAPPER | WeightsMapper(
1518
orig_to_new_prefix={
1519
# Gemma4ForConditionalGeneration already loads the text stack
1520
# from `model.language_model.*`. We reuse that same checkpoint
vllm/model_executor/models/gemma4_mm.py
+6
-3
@@ -40,7 +40,10 @@ from vllm.inputs import MultiModalDataDict
40
from vllm.logger import init_logger
41
from vllm.model_executor.layers.layernorm import RMSNorm
42
from vllm.model_executor.layers.linear import ReplicatedLinear
43
-from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM
43
+from vllm.model_executor.models.gemma4 import (
44
+ _GEMMA4_EXPERT_PARENT_MAPPER,
45
+ Gemma4ForCausalLM,
46
+)
47
from vllm.model_executor.models.module_mapping import MultiModelKeys
48
from vllm.model_executor.models.transformers.utils import recursive_replace_linear
49
from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -998,7 +1001,7 @@ class Gemma4ForConditionalGeneration(
1001
}
1002
1003
# Maps checkpoint prefixes to vLLM module paths.
1001
- hf_to_vllm_mapper = WeightsMapper(
1004
+ hf_to_vllm_mapper = _GEMMA4_EXPERT_PARENT_MAPPER | WeightsMapper(
1005
orig_to_new_prefix={
1006
# vision tower
1007
"model.vision_tower": "vision_tower",
@@ -1010,7 +1013,7 @@ class Gemma4ForConditionalGeneration(
1013
"model.language_model.": "language_model.model.",
1014
"lm_head.": "language_model.lm_head.",
1015
"model": "language_model.model",
1013
- }
1016
+ },
1017
)
1018
1019
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):