Prhub

#32440 fix(gemma4): quantize MTP bridge projections

原始 PR 作者 ayush1399 合并时间 2026-08-20 03:41 文件变更 1 提交数 1 评论 4 代码增减 +2 / -2

执行摘要

修复 Gemma4 MTP 量化配置遗漏,FP8 接受率 0% 提升至 60%

PR body 明确指出:block-FP8 的 gemma-4-31B-it-assistant 检查点 draft-token 接受率为 0%,而 BF16 assistant 约为 65%。检查点将 pre_projection 与 post_projection 存为 block-FP8 权重并带 weight_scale 张量,但这些层以 quant_config=None 构造,导致 FP8 权重加载时缺少 block scales,破坏了 MTP 激活。作者因此将 draft model 的 quant_config 透传给这两个桥接投影层。

值得快速精读。虽然代码改动仅 2 行,但背后是一个完整的调试故事:用 draft-token 接受率作为量化健康的观测指标,反推出 quant_config 传递遗漏。建议关注两点:一是量化配置在不同模型组件间传播时的一致性校验;二是为 pre_projection / post_projection 的量化加载补充专门的回归测试,避免同类问题复发。

讨论亮点

该 PR 几乎没有实质 review 讨论:kpham-sgl 直接 APPROVED,唯一的审查动作是通过 /rerun-test 触发两个 MTP 专项测试并全部通过;PR 作者曾请求 @Jiminator 或 @JustinTong0323 参与 review,但无人留下技术评论。gemini-code-assist[bot] 仅声明其代码审查服务已停止。值得注意的是 checklist 中“添加单元测试”未被勾选,且审查过程无人对此提出异议,说明该修复目前依赖现有 spec 测试作为回归防线。

实现拆解

  1. 定位缺陷:在 python/sglang/srt/models/gemma4_mtp.py 的 MTP 模型 init 中,pre_projection 与 post_projection 两个桥接线性层构造时硬编码 quant_config=None,与主干 Gemma4TextModel 透传 quant_config 的行为不一致。
  2. 修复方式:将两处 quant_config=None 改为 quant_config=quant_config,让桥接投影层继承 draft model 的量化配置,从而在加载 block-FP8 检查点时正确消费 weight_scale 与分块量化参数。
  3. 影响分析:BF16 路径不受影响(quant_config 仍为 None);FP8 路径的 MTP draft-token 接受率由 0% 恢复至约 60%。
  4. 验证与配套:未新增单元测试;通过 /rerun-test 重跑 registered/spec/test_frozen_kv_mtp.py(1-gpu-h100)与 registered/spec/test_gemma4_mtp_31b_extra.py(2-gpu-h100)均通过。
文件 模块 状态 重要度
python/sglang/srt/models/gemma4_mtp.py MTP 模型 modified 5.68

关键符号

Gemma4MTP.__init__

关键源码片段

python/sglang/srt/models/gemma4_mtp.py data-contract

唯一的变更文件,修复 Gemma4 MTP Frozen-KV 桥接投影层 pre_projection / post_projection 的量化配置传递,使 block-FP8 检查点的 weight_scale 能正确加载,将 FP8 draft-token 接受率从 0% 提升至 60%。

def __init__(
    self,
    config: PretrainedConfig,
    quant_config: Optional[QuantizationConfig] = None,
    prefix: str = "",
) -> None:
    text_config = copy.deepcopy(_get_text_config(config))
    text_config.num_kv_shared_layers = 0
    PreTrainedModel.__init__(self, config=text_config)
    self.assistant_config = config
    self.config = text_config
    self.quant_config = quant_config
    self.pp_group = get_pp_group()
​
    self.vocab_size = text_config.vocab_size
    self.hidden_size = text_config.hidden_size
    self.backbone_hidden_size = config.backbone_hidden_size
    self.target_embed_scale = self.backbone_hidden_size ** 0.5
    self.use_ordered_embeddings = bool(
        getattr(config, "use_ordered_embeddings", False)
    )
    self.centroid_intermediate_top_k = int(
        getattr(config, "centroid_intermediate_top_k", 32)
    )
​
    self.target_embed_weight: Optional[torch.Tensor] = None
​
    # 关键修复:pre_projection / post_projection 之前硬编码 quant_config=None,
    # 导致 block-FP8 分片权重加载时丢失 weight_scale,MTP 激活被破坏,
    # draft-token 接受率从约 65%(BF16)跌到 0%(FP8)。
    # 现在透传 draft model 的 quant_config;BF16 路径不受影响(其值仍为 None)。
    self.pre_projection = ReplicatedLinear(
        2 * self.backbone_hidden_size,
        self.hidden_size,
        bias=False,
        quant_config=quant_config,
        prefix=add_prefix("pre_projection", prefix),
    )
    self.model = Gemma4TextModel(
        config=text_config,
        quant_config=quant_config,
        prefix=add_prefix("model", prefix),
    )
    self.post_projection = ReplicatedLinear(
        self.hidden_size,
        self.backbone_hidden_size,
        bias=False,
        quant_config=quant_config,
        prefix=add_prefix("post_projection", prefix),
    )
​
    if text_config.tie_word_embeddings:
        self.lm_head = self.model.embed_tokens
    else:
        self.lm_head = nn.Linear(self.hidden_size, self.vocab_size, bias=False)
    self.logits_processor = LogitsProcessor(text_config, skip_all_gather=True)

评论区精华

MTP 专项测试重跑验证 测试

kpham-sgl 用 /rerun-test 触发 registered/spec/test_frozen_kv_mtp.py(1-gpu-h100)与 registered/spec/test_gemma4_mtp_31b_extra.py(2-gpu-h100),github-actions 返回两个测试全部通过。

结论:与 Gemma4 MTP 路径相关的 spec 测试通过,为合并提供了验证支撑。 · 已解决

缺少针对量化配置传播的单元测试 测试

PR checklist 未勾选“添加单元测试”,2 行修复未附带回归测试;review 过程中无人提出异议。

结论:未补充单元测试,依赖现有 spec 测试间接覆盖;该缺口留待后续补强。 · 待处理

风险与影响

  1. 缺少量化配置传播的单元测试:本 PR 的 2 行修复没有附带回归测试,未来若再次出现类似的 quant_config 透传遗漏,难以被快速发现。
  2. 兼容性风险:如果某个检查点的 pre_projection / post_projection 实际未量化,但传入非 None 的 quant_config,可能导致量化加载路径误判或权重解析报错;当前仅以 gemma-4-31B-it-assistant 的实测结果作为验证。
  3. 影响面受控:修改仅发生在 Gemma4 MTP(Frozen-KV)模型构造路径,对 SRT 其他模型和普通解码路径无影响。

对用户而言,运行 block-FP8 gemma-4-31B-it-assistant 的 speculative decoding 用户获得显著收益:draft-token 接受率从 0% 恢复至约 60%,意味着 MTP 预测从完全失效变为可用,端到端解码吞吐随之改善。对团队而言,该修复暴露了模型层量化配置传播链路的一个系统性风险点,提示后续需要更系统的量化配置一致性检查。整体影响范围限定在 Gemma4 MTP 单一路径,影响程度中等。

量化配置传播遗漏 缺少测试覆盖 核心推理路径变更

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论