#6836 [megatron] fix: return 3-tuple under calculate_per_token_loss to fix MoE aux/z-loss grad blowup at CP>1
原始 PR · 作者 EricMarcus-ai · 合并时间 2026-07-13 10:57
修复 Megatron CP>1 时 MoE aux/z-loss 梯度爆炸
强烈建议使用 Megatron MoE 且 CP>1 的用户合入此 PR。值得关注的设计决策:精确计算路由 token 数的 `_routed_num_tokens` 方法、通过预乘后除抵消的二阶段归一化策略、以及对不合理配置的主动防护。
参与讨论