用 Megatron Core 做逐位确定的预训练
来自 NVIDIA Developer · 期号 2026-10-07
摘要
大模型训练到万亿参数规模时，一次 loss 尖峰可能浪费几千 GPU 天。问题在于训练不是确定性的：同样的代码跑两次，结果可能不一样，没法可靠地复现和调试。 NVIDIA Megatron Core 的方案是给每个 rank 的每一步 tensor 打指纹，离线比对。先确认种子、数据顺序、并行策略一致，然后用渐进式追踪从训练迭代一路定位到具体 kernel。 代价是性能。优化前确定性训练比非确定性慢两位数百分点。NVIDIA 改了 grouped-GEMM kernel：原来多个 writer 写同一个地址，结果取决于到达顺序；改成每个 writer 写私有槽位，最后按固定顺序合并。 三个 Nemotron 负载优化下来，2432 块 GPU 上确定性训练的开销降到约 2%，且 800 步内保持逐位一致。 验证方法也值得注意：光看 loss 不够，两个 run 的 loss 四舍五入后一样，低位可能已经不同。要用 tensor 指纹逐步对比。
深读
导读 大模型训练到万亿参数规模时，一次 loss 尖峰可能浪费数千 GPU 天。问题在于训练不是确定性的：同样的代码跑两次，结果可能不一样，没法可靠地复现和调试。NVIDIA 在 Megatron Core 中给出了一个工程化方案，通过逐位确定性（bitwise determinism）让训练可复现，并将性能开销从两位数百分点压到约 2%。 关键信息 - NVIDIA 在 2,432 块 GPU 上对万亿参数 Nemotron 模型实现了 800 步内的逐位确定性训练，稳态开销约 2%。 - 确定性训练的性能开销（determinism tax）定义为 (确定性步时间 / 基线步时间 - 1) × 100%。 - 优化前，确定性训练比非确定性慢两位数百分点；优化后，三个 Nemotron 负载的开销均降至低个位数。 - 核心 kernel 优化：grouped-GEMM 的多个 writer 原本写入同一地址，改为每个 writer 写私有槽位，最后按固定顺序合并，在保持确定性的同时保留了并行性。 - 验证方法：使用 torch.hash_tensor 对每步的 loss、梯度范数等指标打指纹，离线比对；仅靠打印的 loss 值不够敏感，两个 run 的 loss 四舍五入后可能相同但底层数值已不同。 逐节讲解 什么是逐位确定性，为什么需要它 逐位确定性要求：在数据顺序、架构、并行策略、软件、运行时设置和硬件都固定的前提下，独立启动的两个 run 以及中断后恢复的 run，在每一步都保持逐位一致。Megatron Core 追求两个目标：独立可复现性和检查点恢复一致性。其价值体现在四个方面：可靠复现 loss 尖峰并隔离原因；中断后恢复训练不改变数值轨迹；通过比对发现检查点损坏或软硬件问题；区分有意改动与 run-to-run 随机波动。 如何验证确定性 打印的 loss 值不够敏感，两个 run 的 loss 四舍五入后可能相同，但梯度或参数的低位比特已经不同。更强的做法是：对每步的训练 loss、语言建模 loss、负载均衡 loss、多 token 预测 loss 和梯度范数打指纹，两个 run 逐步比对，第一个不匹配的步数就是缩小调查范围的起点。验证检查点恢复时，需要将不间断的参考 run 与恢复后的 run 比对，确认参数、优化器状态、RNG 状态、数据位置和后续输出都逐位一致。 如何定位和修复确定性破坏 Megatron-LM PR #7262 提出的工作流是：记录每个 rank 的有序 tensor 指纹，离线比对。先确认种子、数据顺序、并行策略、容器和软件栈一致。然后通过渐进式追踪缩小范围：从端到端指标到训练迭代、阶段、模块、操作，最后到 kernel。每个 rank 的指纹写入独立文件，不引入 collectives 或跨 rank 排序，…