使用 NVIDIA Megatron Core 扩展按位确定性预训练

中文精读

原文约 3,438 词 · 20 节

1. 文章信息与导航

   • 这是 NVIDIA 技术博客的一篇关于在 Megatron Core 中实现位级确定性（bitwise determinism）预训练的文章，发布于 2026 年 10 月 6 日。
   • 作者为 Ashwath Aithal、Eric Harper、Santosh Bhavani、Yan Xu 和 Zhiyu Li，归类于 Agentic AI / Generative AI 方向。
   • 页面其余部分为分享按钮、AI 生成摘要、下一步指引、评论入口与标签等导航性内容，正文从简介段开始。

2. 简介：位级确定性为何重要

   • 位级确定性让大规模预训练更易调试、验证与可复现地恢复；在数千 GPU 上训练万亿参数模型时，多并行维度、低精度计算和分布式 checkpoint 让故障复现与修复验证变难，价值尤为突出。
   • 生产和主力训练任务成本高昂，需要较高的成功保证；位级确定性支持可靠的失败重放、loss 尖峰调试、系统变更验证以及中断训练恢复。
   • 在万亿参数规模下，即便很小的变慢也会耗费数千 GPU 天，因此确定性训练必须保持高效。
   • 文章以一个万亿参数 Nemotron 模型为案例，说明 NVIDIA 如何在 Megatron Core 中降低位级确定性预训练的开销。

3. 什么是位级确定性

   • 当数据顺序、架构、配方、并行方式、软件、运行时设置和硬件固定时，位级确定性要求独立运行与 checkpoint 恢复运行走同一条数值轨迹。
   • Megatron Core 针对两项保证，第一项是独立可复现性（independent reproducibility）：从同一初始状态启动的两次运行必须在每个训练步都逐位相同。
   • 第二项保证是位级 checkpoint 恢复（bitwise checkpoint resume）：保存并恢复一个或多个 checkpoint 的运行，必须与不中断的运行完全一致。
   • 图 1 用三条 50 步轨迹展示：独立运行 B 在第 25 步前与连续运行 A 一致；运行 C 在第 25 步保存、第 26 步恢复，到第 50 步仍与 A 一致，两次指纹比较均报告匹配。

4. 为什么预训练需要确定性

   • 位级确定性为大规模预训练提供四项益处。第一项是重放故障（replay failures）：在同一步复现 loss 尖峰，再逐个改动组件以隔离原因。
   • 第二项是保持 checkpoint 连续性（preserve checkpoint continuity）：中断后恢复训练而不改变数值轨迹。
   • 第三项是检测意外差异（detect unexpected differences）：利用原本相同的运行之间出现的不一致，排查 checkpoint 损坏或软硬件问题。
   • 第四项是比较变更（compare changes）：把有意改动的数值影响与运行间波动区分开。
   • 图 2 显示两次运行在第 376 步 language-modeling loss 与零计数一致，第 377 步 language-modeling loss、梯度范数、零计数和参数范数出现差异（370–381 步的 loss 曲线为示意）。

5. 如何验证确定性

   • 打印的 loss 精度不够，两次运行可能打印出相同舍入值，却在梯度或参数的低位比特上不同。
   • 更强的验证是在每一步对选定数值指标取指纹，例如训练 loss、language-modeling loss、负载均衡 loss、多 token 预测（MTP）loss 和梯度范数，两次运行逐步对比，首个不一致步有助于缩小排查范围。
   • 验证位级 checkpoint 恢复时，将不中断的参考运行与恢复运行对比；每次恢复后检查参数、优化器状态、RNG 状态、数据位置及后续输出是否与参考逐位相同。
   • Nemotron 的 checkpoint 路径保留全精度真值源，并用与训练相同的转换路径重建运行时表示；图 3、图 4 分别给出 50 步的 loss/梯度指纹对比与恢复运行的零差异结果。

6. 确定性被破坏时如何修复

   • 在大规模下，非确定性可能是间歇性的或依赖拓扑，可见的 loss 偏离可能远晚于第一个不同比特出现。
   • Megatron-LM PR #7262 提出的工作流会记录按顺序、按 rank 的张量指纹，供离线比较。
   • 第一步是确认种子、数据顺序、batch size、并行方式、容器和软件栈完全相同。

7. 用细粒度追踪定位分歧

   • 先做端到端比较，再把捕获范围从训练迭代逐级缩小到阶段、模块、操作和内核；若两次运行以逐位相同的输入进入某范围、却以不同输出离开，分歧就源于该范围内。
   • 端到端指标负责发现失败并指出第一个序列化指标不同的迭代；宽泛语义追踪覆盖集合通信、流水线通信、重计算、优化器操作和梯度归约，用于识别分歧的训练阶段。
   • 模块与层追踪把分歧缩小到特定模型组件、Transformer 层、MoE 块或优化器阶段；操作级追踪对 ATen 操作和扩展调用取指纹，找出首个输入相同而输出不同的操作；内核级工具确定具体内核、算法配置与导致非确定性的机制。
   • 不要只追踪 loss 明显分离的那次迭代，第一个不同比特可能更早出现：先用宽泛追踪定位最早的迭代与 rank，再只在该范围内开启操作级和内核级追踪。
   • 图 5 以倒漏斗形式展示五个追踪层级，追踪越细、捕获成本越高，范围越窄。

8. 离线比较追踪记录

   • 每个选定 rank 把追踪写入文件，不增加集合通信或跨 rank 排序，以降低掩盖所调查竞态的风险。
   • 按与运行无关的身份对齐事件，例如操作名、出现次数和模块作用域，寻找输入指纹相同而输出指纹不同的情形（见下方代码）。
   • “首个不一致”是各 rank 各自而言，序列号无法建立跨 rank 的顺序。可按因果角色分类：输入相同输出不同→候选来源；输入输出都不同→下游接收者，继续向上游追踪；追踪的 rank 上无来源→来源在捕获范围之外，需扩大迭代窗口或 rank 集合。
   • 若许多 rank 指向同一个源操作，该操作本身很可能非确定；若只有部分 rank 如此，则应排查拓扑、rank 排布、输入分布和归约顺序。

9. 在 GPU 上计算张量指纹

   • 该工作流使用 `torch.hash_tensor` 生成驻留 GPU 的指纹，记录每个张量的形状、dtype、元素数量及摘要。
   • 对 MXFP8 或 NVFP4 张量，要同时对编码值和其 scale 缓冲区取指纹。
   • 整张量 XOR 指纹无法检测置换；对路由图、MoE 分发输出等顺序敏感数据，应用 dim 参数按行或分块取指纹。
   • 指纹是高效检查，但匹配不保证张量逐位相同，需用字节级比较确认逐位相等。

10. 排除误报

    • 调度器盲区：`TorchDispatchMode` 只能观测经由 PyTorch 调度器路由的 `ATen` 操作，自定义内核可能逃出追踪范围；若首个不一致出现在简单的 view、切片或加法上，应探查产生其输入的自定义内核。
    • 探针伪影：未初始化张量、未完成的异步集合通信和非阻塞拷贝可能在内容有效前被读取，须先排除这些情况，再判定生产者非确定。

11. 验证修复

    • 用配对运行检查补丁：原始的一对运行应复现分歧，打补丁后的一对运行应保持逐位相同。
    • 还需确定修复是纠正了非确定性实现，还是绕开了该执行路径，并测量新的性能开销。
    • 位级确定性只在同一硬件与软件环境内验证；跨 GPU 代际、网络配置或库版本的比较不在验证范围内，可能得到不同结果。

12. 为万亿参数 Nemotron 优化确定性训练

    • 有较大开销的确定性配方可用于调试，但用于生产训练代价过高，因此应从开始就协同优化确定性与非确定性执行。
    • 该工作使用的万亿参数 Nemotron 模型同时包含 Mamba 风格的状态空间模型（SSM）层和 Transformer 注意力层。
    • 确定性训练必须覆盖 SSM 与注意力内核、低精度计算、分布式通信以及不同层类型之间的过渡。

13. 建立受控基线

    • 在相同模型、硬件、batch size、并行策略、精度格式、软件环境和测量窗口下，将确定性执行与最快受支持的非确定性配方对比。
    • 测量吞吐与步时间、峰值内存与 GPU 利用率、暴露的通信时间，以及主要内核组的耗时。
    • 据此计算“确定性税”（determinism tax），且只应在训练性能稳定后采集数据。

14. 优化历程

    • 三条优化路线把步时间确定性开销从两位数基线降到低个位数（见图 6，黄色框标出初始差距）。
    • Nemotron 3 Ultra 经 buffer-fill 优化从约 17% 降到 2%，在 3,072 GPU 上进一步降到 1.5%；混合 Triton 代理在恢复 MoE-MLP 融合后从约 60% 降到 37%，再降到 2%。
    • CuTeDSL 优化后的权重梯度路径在 256 GPU 上测得约 3.6% 开销；大规模 Nemotron 配方在 2,432 GPU 上测得约 2% 稳态确定性开销，并在 800 步内保持位级确定性。

15. 内核级优化示例

    • 在 grouped-GEMM 的 epilogue 中，多个 N-tile 原先累加进同一个 `dprob[token]` 地址，结果取决于到达顺序；将这些写者串行化能恢复确定性，但降低并行度。
    • 解决办法是让每个 N-tile 拥有私有输出槽，保留内核内并行执行，并在所有写者完成后按固定顺序合并各槽。
    • 图 7 对比了无序原子加基线、串行写者和并行单写者加固定顺序归约三种时间线，最后一种既逐位精确，又比串行保留更多并行。

16. 在规模上验证优化

    • 少量 GPU 上的低确定性税不保证生产规模同样如此，通信、同步、流水线气泡、专家路由和负载均衡都会改变相对开销。
    • 举例说明：假设一个在 10,000 GPU 上非确定性基线为 100 天的训练，把开销从 15% 降到 5% 可节省 10 天，即 100,000 GPU 天。

17. 随模型、内核和配方演进保持确定性

    • 新内核、融合、精度格式或并行配置都可能引入非确定性，Megatron-LM 用三类检查防范回归。
    • 配方验证：`--deterministic-mode` 应用规范化环境设置、启用 PyTorch 确定性算法，并拒绝没有确定性路径的特性。
    • 内核测试：用相同输入和 RNG 状态重复操作，逐字节比较输出与梯度；模块验证：在恢复 RNG 状态后跨并行配置（含 FP8 与 FP4）重复模型和 Transformer 块，再用独立端到端运行确认全精度训练指标逐位相同。
    • 当前覆盖范围与已知缺口可查 determinism status、operation catalog 和 kernel testing guide 三个文档页面。

18. 非确定性的常见来源

    • 表 2 按症状给出原因、诊断与解决方式：从第一步起就发散但不稳定，多为运行时自动调优选中了不同内核配置，应比较所选配置并追踪最早受影响的输出，然后固定或缓存已验证配置。
    • MoE 运行仅在启用融合辅助 loss 时发散，是归约顺序未固定，应逐个关闭融合追踪 loss 计算，改用确定性归约或禁用该融合。
    • 恢复运行逐渐偏离连续运行，是 RNG 状态恢复不完整，应比较恢复前后的 RNG 状态，逐位保存和恢复每条 RNG 流。
    • 软件或容器更新后出现分歧，多为底层库内核变化，应二分排查软件栈并构建最小内核复现，在修正前选用确定性内核路径；加载 checkpoint 后参数立即不同，多为低精度权重或 scale 重建方式不同，应保留全精度真值源与精确重建路径。
    • 小规模通过而大规模失败的配方，多与规模相关的通信或专家并行路径有关，应系统增大规模并追踪选定 rank，隔离并修正第一个随规模变化的操作。

19. 开始使用

    • 位级确定性使大规模预训练更易调试和恢复，Nemotron 案例说明内核与配方优化能降低其性能代价；应计划使用的硬件和软件环境中验证可复现性与开销。
    • 入门方法是查阅 Megatron Core 用户指南中的受支持配方，然后对比独立运行与 checkpoint 恢复运行以验证可复现性。
    • 文末标签涉及 Agentic AI/生成式 AI、数据中心/云、Nemotron、中级技术、最佳实践、featured 与 Megatron。

20. 关于作者

    • Ashwath Aithal 是 NVIDIA 的 AI 与深度学习系统和算法总监，领导为大型 AI 模型的训练与优化开发软件的团队；工作涵盖 NVIDIA Megatron Core、NVIDIA NeMo 和 NeMo RL，重点是可扩展的分布式训练、强化学习与性能优化。他拥有罗格斯大学电气与计算机工程硕士学位。
    • Eric Harper 是 NVIDIA 的深度学习工程经理，负责 Megatron Core 的工程工作，该库是用于大规模训练大语言模型的开源库；其团队聚焦提升 NVIDIA GPU 上分布式训练软件的性能、可扩展性与易用性。他拥有迈阿密大学数学博士学位。
    • Santosh Bhavani 是 NVIDIA 的产品经理，从事深度学习框架、Megatron Core 和 Transformer Engine 相关工作，拥有卡内基梅隆大学计算机科学学士学位。
    • Yan Xu 是 NVIDIA 的高级 AI 软件工程师，在 Megatron Core 上从事大规模语言模型训练，开发使预训练更高效、可复现、更易调试的工具与技术；拥有马里兰大学帕克分校计算机科学硕士学位。
    • Zhiyu Li 是 NVIDIA 的高级深度学习工程师，专长于高效、高性能的大语言模型（LLM）训练系统与基础设施，包括预训练与后训练框架；拥有卡内基梅隆大学计算机科学硕士学位。