#7702 [megatron] fix: pad VLM THD input to TP/CP alignment before SP scatter
原始 PR · 作者 HollowMan6 · 合并时间 2026-09-11 10:17
对齐 VLM THD 序列长度,修复 Megatron SP 崩溃
建议精读,尤其是正在维护 Megatron VLM / 序列并行 / 上下文并行前向路径的工程师。值得关注的三个设计点:一是「SP scatter 要求 S 能被 tp_size 整除」这一硬约束通过 padding 在模型入口处满足,而不是在 scatter 处打补丁;二是 CP>1 时对齐粒度取 `tp_size * cp_size * 2` 的推导(zigzag 切分保证单 rank 局部长度仍可被 tp_size 整除);三是把对齐逻辑收敛到 `build_vlm_attn_mask_thd`、让 fused 路径复用,避免两套实现继续分叉。同时应把 Copilot 的测试建议当作待办跟进:补一个覆盖 B=1、4026 这类奇数长度的 CPU 单测,断言输出 S 可被 `tp_size`(CP>1 时为 `tp_size * cp_size * 2`)整除,并在 CP 分支校验每个 rank 拿到完整切片。
参与讨论