Prhub

#21347 [Bugfix] Fix PP tied embeddings weight loading for qwen3.5 4B dense model

原始 PR 作者 edwingao28 合并时间 2026-04-01 14:51 文件变更 1 提交数 4 评论 19 代码增减 +22 / -0

执行摘要

修复 Qwen3.5 4B dense 模型在 PP=2 时权重加载错误导致的输出乱码问题。

修复 issue #21093 中报告的 Qwen3.5-4B dense 模型在 PP>1 时输出不正确问题。PR body 指出根因是 Qwen3_5ForConditionalGenerationQwen3_5MoeForConditionalGeneration 覆盖了父类 Qwen3VLForConditionalGenerationload_weights 方法但省略了 tie_word_embeddings 处理,在 PP>1 时 lm_head 作为单独的 ParallelLMHead 缺少独立权重,未复制 embed_tokens.weight 导致输出乱码。

该 PR 值得精读,特别是关注模型权重加载机制和 PP 下的初始化设计决策,有助于理解大型语言模型在分布式环境中的权重处理模式。

讨论亮点

Review 中主要有两个讨论点:1) ShangmingCai 建议先检查 self.config.tie_word_embeddings 以实现早期退出,优化性能,作者 edwingao28 采纳并更新代码;2) yuan-luo 指出 MoE 模型的 load_weights 方法也可能需要相同修复,作者回应当前 MoE 检查点 tie_word_embeddings 为 false,但添加了防御性守卫以确保未来兼容性。讨论聚焦于正确性和设计优化。

实现拆解

修改集中在文件 python/sglang/srt/models/qwen3_5.pyload_weightsload_fused_expert_weights 方法中。添加了条件逻辑:当 self.config.tie_word_embeddings 为 true、当前 rank 是最后一个(self.pp_group.is_last_rank)且加载的权重名称为 "model.embed_tokens.weight" 时,将权重复制到 lm_head.weight。这匹配了父类 Qwen3VLForConditionalGeneration 的权重加载模式,确保了 lm_head 正确初始化。

文件 模块 状态 重要度
python/sglang/srt/models/qwen3_5.py models/qwen3_5 modified 8.0

分析完成后,这里会展示 LLM 生成的相对完整源码片段和详细注释。

关键符号

load_weights load_fused_expert_weights

评论区精华

检查顺序优化 正确性

ShangmingCai 建议先检查 `self.config.tie_word_embeddings` 以实现早期退出,优化性能。

结论:作者采纳建议,更新代码以确保更高效的条件判断。 · 已解决

MoE 模型扩展 正确性

yuan-luo 指出 Qwen3_5MoeForConditionalGeneration 的 load_weights 方法也可能需要相同修复。

结论:作者回应当前 MoE 检查点 tie_word_embeddings 为 false,但添加了防御性守卫以确保未来兼容性。 · 已解决

风险与影响

风险较低,因为修改是添加缺失的逻辑,与父类模式保持一致。但需注意:MoE 模型虽当前无问题,防御性守卫处理了边缘情况;缺少单元测试可能掩盖回归问题,但 PR 提供了准确性验证(如 gsm8k 测试),降低了风险。

对用户而言,修复了模型在 PP 部署下的输出正确性,避免了垃圾输出,提升了使用体验。系统层面,确保了权重加载的一致性,影响范围限于使用 Qwen3.5 dense 模型且 PP>1 的场景,对整体系统稳定性有积极影响。

缺少测试覆盖 边缘情况依赖现有检查点

关联 Issue

#21093 [Bug] [Qwen3.5 Dense] pp=2 Qwen3.4-4B result is incorrect

完整报告

执行摘要

本 PR 修复了 Qwen3.5-4B dense 模型在并行处理(PP)大于 1 时因权重加载错误导致的输出乱码问题,通过添加缺失的 tie_word_embeddings 处理逻辑,确保了模型在分布式环境下的正确性。

功能与动机

动机源于 issue #21093,用户报告 Qwen3.5-4B dense 模型在 PP=2 时输出不正确。根因是模型类 Qwen3_5ForConditionalGenerationQwen3_5MoeForConditionalGeneration 覆盖了父类 Qwen3VLForConditionalGenerationload_weights 方法但遗漏了 tie_word_embeddings 权重复制,导致在 PP>1 时 lm_head 权重未初始化,直接产生垃圾输出。

实现拆解

修改集中在 python/sglang/srt/models/qwen3_5.py 文件。在 load_weightsload_fused_expert_weights 方法中,添加了以下逻辑:

if (
    self.config.tie_word_embeddings
    and self.pp_group.is_last_rank
    and "model.embed_tokens.weight" in name
):
    if "lm_head.weight" in params_dict:
        lm_head_param = params_dict["lm_head.weight"]
        weight_loader = getattr(lm_head_param, "weight_loader", default_weight_loader)
        weight_loader(lm_head_param, loaded_weight)

这确保了当 tie_word_embeddings 为 true 时,最后一个 rank 的 lm_head 权重从 embed_tokens 复制,与父类模式一致。

评论区精华

Review 中主要讨论点:

  • 检查顺序优化:ShangmingCai 建议“先检查 self.config.tie_word_embeddings 以实现早期退出”,作者采纳并更新代码,提升了条件判断效率。
  • MoE 模型扩展:yuan-luo 指出“MoE 模型的 load_weights 也可能需要相同修复”,作者回应当前 MoE 检查点 tie_word_embeddings 为 false,但添加了防御性守卫以确保未来兼容性。

风险与影响

风险较低:修改逻辑与父类一致,但缺少单元测试可能隐藏回归问题。MoE 模型的防御性守卫处理了边缘情况,降低了风险。影响方面,修复了模型在 PP 部署下的输出正确性,提升了用户体验,影响范围限于使用 Qwen3.5 dense 模型且 PP>1 的场景。

关联脉络

与此 PR 相关的历史 PR 包括 #17122(同为模型 bugfix,展示了类似的设计模式),以及 issue 中提及的 PR #21070(修复了同一模型的 PP 分割 OOM 问题,但未解决权重加载)。这表明在大型语言模型集成中,权重加载和分布式处理是需要持续关注的关键区域。

参与讨论