Skip to content

fix: correct final hidden normalization for MTP - #1614

Open
sufubao wants to merge 9 commits into
ModelTC:mainfrom
sufubao:fix-qwen35-draft-norm
Open

sufubao wants to merge 9 commits into
ModelTC:mainfrom
sufubao:fix-qwen35-draft-norm

Conversation

@sufubao

@sufubao sufubao commented Oct 8, 2026 •

Copy link
Copy Markdown
Collaborator

修复 MTP hidden 的 final norm 约定:Qwen3.5 和 GLM5 Next 在后续 draft 步骤中固定套用 main norm;DeepSeek 的 main/draft、GLM4 MoE Lite 的 main 则漏传了归一化后的 hidden。公共输出层按 FinalHiddenPolicy 枚举准备 NONE、PRE_NORM 或 POST_NORM,norm 与投影分开,logits 和需要传递的归一化 hidden 共用一次计算。普通非 MTP prefill 仍先选最后位置再做 norm。

逐一核对了全部已注册 draft 家族,保留各模型的差异:

模型 main 最终 hidden draft 最终 hidden
Qwen3.5 Dense / MoE POST_NORM POST_NORM
GLM5 Next POST_NORM POST_NORM
DeepSeek POST_NORM POST_NORM
GLM4 MoE Lite POST_NORM PRE_NORM
Mistral PRE_NORM PRE_NORM,输入侧使用共享 final norm
Qwen3 MoE PRE_NORM PRE_NORM,保留现有约定
EAGLE3 中间层 hidden PRE_NORM
DFlash / DFlash2 / DSpark,包括 Qwen3.5 adapter 中间层 hidden 不输出最终 hidden

Qwen3.5、GLM5 Next 删除 draft 输入端固定的 main norm 调用与绑定。BaseModel 直接收集 PostLayerOutput.final_hidden,漏返回不再回退到原始 hidden;reward head 复用公共准备逻辑,DSpark 直接投影路径显式执行 norm。

参考 vLLM 的 GLM5 Next MTP、DeepSeek MTP、GLM4 MoE Lite main 和 draft。GLM4 MoE Lite draft 按其参考实现保留 PRE_NORM,不随主模型继承 POST_NORM。

MTP 输出约定统一声明在 draft registry 的 MtpPairConfig 中,按 draft checkpoint 的 model_type 与 MTP 模式选择,分别配置 main_final_hidden_policy 和 draft_final_hidden_policy。模型类不再持有 mtp_final_hidden_policy,避免 draft 继承 main 策略。MtpManager 为 main 读取首个 draft 配对,为 draft 读取自身 checkpoint 的配置,并缓存解析结果;非 MTP 返回 NONE,不读取 draft 配置。策略枚举和配对配置位于轻量的 common/mtp_config.py,保持模型注册与 tokenizer 的延迟导入。

BaseModel 将有效策略写入 infer_state.final_hidden_policy,post layer 用当前模型的 norm 权重返回 PostLayerOutput.final_hidden;collector 只收集结果,不持有 policy。autotune 使用 NONE。

验证:540 项相关测试全部通过,覆盖全部注册配对的 main/draft 策略、同一 main 切换配对、多个 draft checkpoint、配置缓存、infer-state 传递、CPU FP32/CUDA BF16 连续两步 draft、post layer、DSpark、CUDA graph、proposer、draft registry 与 tokenizer 导入隔离。已合入最新 main,并补齐 DFlash2 的 NONE/NONE 配对策略;本轮完整选集无跳过、无排除。pre-commit(Black、flake8)和 diff 检查通过。测试通过 exp 留档:261014-151837-python-m-pytest-q-unit-tests-common-basemodel-te。本轮未运行完整模型接受率、吞吐或真实 TP/DP 多卡对照。

#1624 分支已同步至本 PR 最新提交 69ff8c0,在此基础上统一 collector 输出为 ModelMtpOutput.spec_hidden;按要求保持关闭。

@sufubao
sufubao requested a review from shihaobai October 8, 2026 09:27
@sufubao sufubao changed the title fix: normalize Qwen3.5 MTP hidden with producer weights fix: correct final hidden normalization for MTP Oct 12, 2026
@sufubao sufubao closed this Oct 14, 2026
@sufubao sufubao reopened this Oct 14, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant