A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 01:15 UTC
[BugFix] Fix Gemma4 'layers.0.moe.experts.0.down_proj_packed' KeyError issue (#40708)
Signed-off-by: SoluMilken <ypiheyn.imm02g@g2.nctu.edu.tw>
SoluMilken committed
May 10, 2026 at 01:20 UTC
25abddc1a5cb5dbe51d629aada9360579acb8acf
2 files changed
+55
-18
tests/models/quantization/test_awq.py
+31
@@ -93,6 +93,37 @@ def run_awq_test(
93
)
94
95
96
+@pytest.mark.parametrize(
97
+ ("model", "quantization", "dtype"),
98
+ [
99
+ ("mattbucci/gemma-4-26B-AWQ", "awq", "float16"),
100
+ ("cyankiwi/gemma-4-26B-A4B-it-AWQ-4bit", "compressed-tensors", "bfloat16"),
101
+ ],
102
+ ids=[
103
+ "gemma4-moe-standard-awq-dot-suffix",
104
+ "gemma4-moe-compressed-tensors-underscore-suffix",
105
+ ],
106
+)
107
+@torch.inference_mode()
108
+def test_awq_load(
109
+ vllm_runner: type[VllmRunner],
110
+ example_prompts: list[str],
111
+ model: str,
112
+ quantization: str,
113
+ dtype: str,
114
+) -> None:
115
+ """Regression test: AWQ weight loading must not KeyError."""
116
+ with vllm_runner(
117
+ model,
118
+ quantization=quantization,
119
+ dtype=dtype,
120
+ max_model_len=128,
121
+ enforce_eager=True,
122
+ ) as vllm_model:
123
+ outputs = vllm_model.generate_greedy(example_prompts[:2], max_tokens=32)
124
+ assert len(outputs) == 2
125
+
126
+
127
@pytest.mark.parametrize(
128
("source_model", "quant_model"),
129
[("OpenGVLab/InternVL2-2B", "OpenGVLab/InternVL2-2B-AWQ")],
vllm/model_executor/models/gemma4.py
+24
-18
@@ -40,6 +40,7 @@ from vllm.model_executor.layers.attention import Attention
40
from vllm.model_executor.layers.fused_moe import (
41
FusedMoE,
42
GateLinear,
43
+ fused_moe_make_expert_params_mapping,
44
)
45
from vllm.model_executor.layers.layernorm import RMSNorm
46
from vllm.model_executor.layers.linear import (
@@ -1368,30 +1369,35 @@ class Gemma4Model(nn.Module, EagleModelMixin):
1369
# moe.experts.{id}.gate_proj → FusedMoE w1 (shard of w13)
1370
# moe.experts.{id}.up_proj → FusedMoE w3 (shard of w13)
1371
# moe.experts.{id}.down_proj → FusedMoE w2
1371
- #
1372
- # Use prefix matching to handle both weights and
1373
- # quantization scale parameters. The param_name is a prefix ending
1374
- # in underscore, and weight_name ends with a dot, so that:
1375
- # "experts.0.gate_proj.weight_scale" -> "experts.w13_weight_scale"
1376
- # "experts.0.gate_proj.weight" -> "experts.w13_weight"
1372
num_experts = getattr(self.config, "num_experts", None) or 0
1378
- expert_params_mapping = [
1379
- # (param_name, weight_name, expert_id, shard_id)
1373
+ # Strategy A: dot-separated suffix
1374
+ # (standard AWQ/GPTQ e.g. .qweight, .scales, .weight)
1375
+ dot_suffix_expert_params_mapping = fused_moe_make_expert_params_mapping(
1376
+ self,
1377
+ ckpt_gate_proj_name="gate_proj",
1378
+ ckpt_down_proj_name="down_proj",
1379
+ ckpt_up_proj_name="up_proj",
1380
+ num_experts=num_experts,
1381
+ )
1382
+ # Strategy B: underscore-separated suffix
1383
+ # (CompressedTensors-format AWQ/W4A16 _packed, _scale)
1384
+ underscore_suffix_expert_params_mapping = [
1385
(
1381
- "experts.w13_"
1382
- if proj_name in ["gate_proj", "up_proj"]
1383
- else "experts.w2_",
1384
- f"experts.{expert_id}.{proj_name}.",
1386
+ f"{param_name}weight_",
1387
+ f"{weight_name.rstrip('.')}_",
1388
expert_id,
1389
shard_id,
1390
)
1388
- for expert_id in range(num_experts)
1389
- for shard_id, proj_name in [
1390
- ("w1", "gate_proj"),
1391
- ("w2", "down_proj"),
1392
- ("w3", "up_proj"),
1393
- ]
1391
+ for (
1392
+ param_name,
1393
+ weight_name,
1394
+ expert_id,
1395
+ shard_id,
1396
+ ) in dot_suffix_expert_params_mapping
1397
]
1398
+ expert_params_mapping = (
1399
+ dot_suffix_expert_params_mapping + underscore_suffix_expert_params_mapping
1400
+ )
1401
params_dict = dict(self.named_parameters())
1402
# Include buffers (e.g. layer_scalar) so they can be loaded too
1403
params_dict.update(dict(self.named_buffers()))