执行摘要

该PR扩展了Gemma4模型的权重加载逻辑，以支持量化MoE（Mixture of Experts）检查点。通过改进专家参数映射和权重迭代器，能够处理每专家2D量化权重和缩放参数，确保量化模型（如FP8动态量化）能够正确加载并保持推理精度。变更集中在单个文件vllm/model_executor/models/gemma4.py，风险较低，但需注意映射逻辑的兼容性。

功能与动机

为什么做：为了支持Gemma4量化MoE模型的加载。PR body中明确指出，目的是“扩展gemma4 MoE权重加载以包含2D量化层和参数的逻辑”。作者提供了量化检查点（如RedHatAI/gemma-4-26B-A4B-it-FP8-Dynamic）的加载示例，并展示了量化模型在GSM8K基准测试上的性能（0.9669 vs 原始0.9702），表明量化模型在精度损失极小的情况下能够加载运行。

实现拆解

主要修改了vllm/model_executor/models/gemma4.py中的两个函数：

1. load_weights函数：

   • 重构了专家参数映射逻辑，从仅支持3D打包张量扩展到同时支持3D和2D量化权重。
   • 使用前缀匹配处理量化参数（如.weight_scale），映射规则示例：
     • "experts.0.gate_proj.weight_scale" → "experts.w13_weight_scale"
     • "experts.0.gate_proj.weight" → "experts.w13_weight"
     • 移除了对加载权重必须为2D的断言，因为量化参数可能是1D或标量。

2. _weight_iterator函数：

   • 添加正则表达式重映射：name = re.sub(r"\.experts\.(\d+)\.", r".moe.experts.\1.", name)
   • 将.experts.{id}.{proj}转换为.moe.experts.{id}.{proj}，以处理每专家2D量化权重。

评论区精华

review中主要讨论了两个技术点：

1. 正则表达式潜在问题：

   gemini-code-assist[bot]：“The regular expression \.experts\.(\d+)\. might lead to incorrect remapping if the weight name already contains the .moe. prefix... Consider using a negative lookbehind.”

kylesayrs：“I don't think that saving vLLM checkpoints is a real use case.”

结论：未采纳建议，认为实际风险低。

1. 映射优先级优化：

   kylesayrs：“It seems like you could also handle this by leaving the current code unchanged and just adding another mapping (with higher priority)”

mgoin：“What is the point if this will always match first?”

结论：采用前缀匹配逻辑，未添加额外映射。

风险与影响

风险：

• 兼容性风险：修改了权重加载逻辑，可能影响非量化MoE检查点的加载。但PR body测试表明原始检查点不受影响。
• 映射错误风险：正则表达式重映射可能在某些边缘情况下产生错误前缀（如.moe.moe.experts），但讨论认为实际风险低。
• 精度风险：量化模型加载后精度略有下降（GSM8K从0.9702降至0.9669），但属于预期范围内的量化损失。

影响：

• 用户：Gemma4量化MoE模型用户现在可以加载和使用量化检查点，扩展了模型部署选项。
• 系统：仅影响Gemma4模型的权重加载路径，对系统其他部分无影响。
• 团队：为后续量化MoE模型支持提供了参考实现。

关联脉络

从近期历史PR看，本PR与以下PR相关：

• PR #39181：修复Qwen系列MoE精度问题，但针对不同模型系列。
• PR #39322：添加NVFP4线性层批量不变性支持，同样涉及量化，但针对不同层类型。

本PR是vLLM对量化MoE模型支持的一部分，反映了团队在扩展量化模型覆盖范围上的持续努力。