返回 文章 build CMS 文章

NVIDIA Megatron Core 实现位级确定性预训练：万亿参数 Nemotron 案例研究

位级确定性让万亿参数预训练可调试、可验证、可复现恢复，且通过优化将性能开销降至个位数。

NVIDIA
位级确定性Megatron Core大规模预训练Nemotron
成长分 / 100 79 综合收获、行动、留存与影响

NVIDIA Megatron Core 实现位级确定性预训练：万亿参数 Nemotron 案例研究
为什么值得读了解位级确定性如何解决大规模预训练中故障复现、检查点连续性和变更验证的难题。

学习 NVIDIA 在万亿参数 Nemotron 模型上降低确定性开销的优化历程和内核级技巧。

关键洞察
  1. 位级确定性要求独立运行和检查点恢复运行在固定条件下遵循相同的数值轨迹，提供独立可复现性和位级检查点恢复两大保证。
  2. 打印损失值不足以验证位级相同，需对训练损失、梯度范数等指标进行逐步指纹识别，并比较首个不匹配步骤。
  3. 渐进式追踪从端到端指标缩小到内核级，通过离线比较按 rank 的 tensor 指纹定位发散源，避免掩盖竞态。
转成行动

深入阅读

正文与原文对照

原文保真覆盖：全文原文字符：17087

位级确定性使大规模预训练更易于调试、验证和可复现地恢复。当在数千个 GPU 上训练数万亿参数的模型时，这些优势变得尤为宝贵，因为多种并行维度、低精度计算和分布式检查点使故障复现和修复验证变得复杂。

生产和英雄训练运行成本高昂，需要强有力的成功保证。位级确定性能够实现可靠的故障重放、损失尖峰调试、系统变更验证以及从中断的训练运行中恢复。在万亿参数规模下，即使很小的减速也可能耗费数千个 GPU 日，因此确定性训练必须保持高效。一个万亿参数的 Nemotron 案例研究展示了 NVIDIA 如何降低 NVIDIA Megatron Core 中位级确定性预训练的开销。

什么是位级确定性？

位级确定性要求，当数据顺序、架构、配方、并行性、软件、运行时设置和硬件固定时，独立运行和检查点恢复的运行遵循相同的数值轨迹。Megatron Core 的目标是实现两个保证。

独立可复现性： 从相同初始状态启动的两次运行必须在每个训练步骤保持位级相同。

位级检查点恢复： 保存并恢复一个或多个检查点的运行必须与不间断运行保持相同。

50 步内的三条训练轨迹。独立运行 B 在第 25 步之前与连续运行 A 匹配。运行 C 在第 25 步保存，在第 26 步恢复，并在第 50 步之前与运行 A 匹配。两个指纹比较均报告匹配。

图 1. 独立运行和检查点恢复的运行复现出一条位级相同的训练轨迹

为什么预训练应该是确定性的？

位级确定性为大规模预训练提供四个好处。

重放故障： 在同一步骤复现损失尖峰，然后一次更改一个组件以隔离原因。保持检查点连续性： 在不改变数值轨迹的情况下恢复中断的训练。检测意外差异： 利用其他方面相同的运行之间的不匹配来调查检查点损坏或软件和硬件问题。比较变更： 将预期变更的数值效应与运行间差异区分开来。

两次运行在第 376 步具有匹配的语言建模损失和零计数。在第 377 步，语言建模损失、梯度范数、零计数和参数范数不同。插图比较了测得的损失值。从第 370 步到第 381 步的扩展损失曲线是示意性的。

图 2. 全精度指标揭示了第 377 步的首次测量不匹配

如何验证确定性

打印的损失值不够敏感，无法确定位级相同的数值。两次运行可能打印相同的舍入损失，但其梯度或参数的低位比特不同。

更强大的验证程序在每个步骤对选定的数值指标进行指纹识别。这些指标可以包括训练损失、语言建模损失、负载均衡损失、多令牌预测（MTP）损失和梯度范数。运行相同的配置两次，并逐步比较这些指纹。第一个不匹配的步骤有助于缩小调查范围。

为了验证按位检查点恢复，请将一次不间断的参考运行与检查点恢复运行进行比较。每次恢复后，验证参数、优化器状态、RNG 状态、数据位置以及后续输出是否与参考运行按位完全相同。

万亿参数 Nemotron 模型的检查点路径保留了一个全精度的真实来源，并通过与训练期间相同的转换路径重建运行时表示。

六组配对运行图展示了语言建模损失、梯度范数、MTP-1 损失、MTP-1 词元加权损失、MTP-2 损失和 MTP-2 词元加权损失。所比较的运行在第 50 步之前重叠。

图 3. 配对运行轨迹比较了 50 个训练步中的损失和梯度指纹

不间断运行和恢复运行的 MTP-2 损失轨迹在 50 步中重叠。全精度差异图保持为零，摘要报告了保存范围、恢复范围、完整轨迹和最大增量的精确匹配。

图 4. 检查点恢复运行在第 25 步检查点前后均与不间断参考运行匹配

当确定性被破坏时如何修复

在大规模场景下，非确定性可能是间歇性的或依赖于拓扑结构，而可见的损失发散可能出现在第一个不同比特之后很久。Megatron-LM PR #7262 中提出的工作流记录了有序的、按 rank 的 tensor 指纹，用于离线比较。首先，确认种子、数据顺序、批量大小、并行度、容器和软件栈完全相同。

通过细粒度追踪定位发散

从端到端比较开始，然后将捕获范围从训练迭代缩小到阶段、模块、操作和内核。如果两次运行以按位相同的输入进入某个范围，但以不同的输出离开该范围，则发散源于该范围内部。

  • 端到端指标检测故障并识别序列化指标首次出现差异的迭代。
  • 广泛的语义追踪覆盖集合通信、流水线通信、重计算、优化器操作和梯度归约，以识别发散的训练阶段。
  • 模块和层追踪将发散缩小到特定的模型组件、transformer 层、专家混合（MoE）块或优化器阶段。
  • 操作级追踪对 ATen 操作和有针对性的扩展调用进行指纹记录，以识别第一个输入匹配但输出不同的操作。
  • 内核级工具识别负责的内核、算法配置和非确定性机制。

不要只追踪损失明显分离的那个迭代；第一个不同的比特可能更早出现。使用广泛追踪定位最早的发散迭代和 rank，然后仅在缩小后的范围内启用详细的操作和内核追踪。

倒漏斗图包含五个追踪级别：端到端指标、广泛语义追踪、模块和层追踪、操作级追踪和内核级工具。随着追踪细节和捕获成本增加，范围逐渐缩小。

图 5. 渐进式追踪将确定性故障从端到端指标缩小到负责的内核

离线比较追踪

每个被选中的 rank 将其 trace 写入文件，不添加集合操作或跨 rank 排序。这降低了掩盖正在调查的竞态的风险。通过一个与运行无关的身份来对齐事件，例如操作名称、出现次数和模块作用域。

寻找匹配的输入指纹和不同的输出指纹。

hash(input_run_A) == hash(input_run_B)
hash(output_run_A) != hash(output_run_B)

“首个不匹配”分别针对每个 rank 而言；序列号无法跨 rank 建立顺序。按因果角色对每个 rank 的首个不匹配进行分类。

某个 rank 上的首个不匹配 解释 下一步操作
输入匹配；输出不同 候选来源 调查此操作或其隐藏的生产者
输入和输出均不同 下游接收者 继续向上游追踪
被追踪的 rank 上无来源 来源位于捕获范围之外 扩大迭代窗口或 rank 集合

表 1. 如何解释某个 rank 上追踪到的首个不匹配

如果许多 rank 识别出同一个起源操作，则该操作本身很可能是非确定性的。如果只有一部分 rank 识别出，则调查拓扑结构、rank 放置、输入分布和归约顺序。

GPU 上的指纹张量

所提出的工作流使用 torch.hash_tensor

来生成驻留 GPU 的指纹。记录每个张量的形状、dtype 和元素数量，以及其摘要。对于 MXFP8 或 NVFP4 张量，对编码值及其缩放缓冲区都进行指纹计算。

全张量 XOR 指纹无法检测排列。对于顺序敏感的数据，如路由映射和 MoE 调度输出，使用 dim 参数按行或块对这些张量进行指纹计算。指纹是一种高效的检查，但匹配并不能保证张量逐位相同。使用字节级比较来确认逐位相等。

排除误报

在确定原因之前，检查追踪盲点和无效读取。

调度器盲点。TorchDispatchMode

观察通过 PyTorch 调度器路由的ATen

操作。自定义内核可能逃逸此追踪范围。如果首个不匹配出现在简单的视图、切片或加法操作上，则探查产生其输入的自定义内核。探查伪影。未初始化的张量、不完整的异步集合通信以及非阻塞拷贝可能在其内容有效之前被读取。在声明生产者非确定性之前，排除这些情况。

验证修复

通过配对运行检查来验证补丁。

  • 原始配对应重现分歧。
  • 打补丁后的配对应保持逐位相同。
  • 确定修复是纠正了非确定性实现，还是绕开了它。
  • 测量新的性能成本。

注意，逐位确定性是在相同的硬件和软件环境中验证的。跨 GPU 代际、网络配置或库版本的比较不在本验证范围内，可能会产生不同的结果。

为万亿参数 Nemotron 模型优化确定性训练

具有大量开销的确定性方案可能支持调试，但对于生产训练来说成本过高。确定性和非确定性执行应从一开始就协同优化。这项工作使用了一个万亿参数的 Nemotron 模型，它结合了 Mamba 风格的状态空间模型（SSM）层和 Transformer 注意力层。确定性训练必须覆盖 SSM 和注意力内核、低精度计算、分布式通信以及层类型之间的转换。

建立受控基线

使用相同的模型、硬件、批量大小、并行策略、精度格式、软件环境和测量窗口，将确定性执行与支持的最快非确定性方案进行比较。

测量：

  • 吞吐量与步长时间
  • 峰值内存与 GPU 利用率
  • 暴露的通信时间
  • 主要内核组中花费的时间

计算确定性税：

仅在训练性能稳定后收集数据。

优化历程

三条优化路径展示了步长时间确定性开销。Nemotron 3 Ultra 在缓冲区填充优化后从约 17% 降至 2%，然后在 3,072 个 GPU 上降至 1.5%。混合 Triton 代理在恢复 MoE-MLP 融合后从约 60% 降至 37%，然后降至 2%。CuTeDSL 路径从确定性中断进展到优化的权重梯度路径，在 256 个 GPU 上约为 3.6%，然后在 2,432 个 GPU 上通过 800 步的位级确定性降至约 2%。黄色框突出显示了初始性能差距。

图 6. Nemotron 3 Ultra 的步长时间确定性开销从约 17% 降至约 1.5%，混合 Triton 代理从约 60% 降至约 2%；CuTeDSL 路径在 2,432 个 GPU 上达到约 2%

跨三个工作负载的优化将步长时间确定性开销从两位数的基线降低到较低的个位数。Nemotron 3 Ultra 在缓冲区填充优化后从约 17% 提高到 2%，然后在 3,072 个 GPU 上提高到 1.5%。混合 Triton 代理在恢复 MoE-MLP 融合后从约 60% 提高到 37%，然后提高到 2%。CuTeDSL 优化的权重梯度路径在 256 个 GPU 上测得约 3.6% 的开销。在 2,432 个 GPU 上，大规模 Nemotron 配方测得约 2% 的稳态确定性开销，在 800 步内具有位级确定性。

内核级优化示例

在分组 GEMM 的尾声（epilogue）中，多个 N 块最初累积到同一个 dprob[token] 地址，使得结果依赖于它们的到达顺序。将这些写入者串行化恢复了确定性，但降低了并行性。

为每个 N 块分配一个私有输出槽，在内核内部保持并行执行，并在所有写入者完成后按固定顺序合并这些槽。

三个分组 GEMM 时间线比较了无序原子添加基线、串行化写入者以及并行单写入者槽后跟固定顺序归约。最终方法在保持比串行化更多并行执行的同时实现了位精确。

图 7. 分离分组 GEMM 写入者并应用固定归约顺序，在保持确定性的同时实现了更多并行性

大规模验证优化

在少量 GPU 上的低确定性税并不能保证在生产规模下获得相同的结果。通信、同步、流水线气泡、专家路由和负载均衡可能会改变相对开销。对于在 10,000 个 GPU 上运行 100 天的非确定性基线的假设运行，将开销从 15% 降低到 5% 将节省 10 天，即 100,000 GPU 天。

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

新的内核、融合、精度格式或并行配置可能会引入非确定性。Megatron-LM 通过以下检查防止回归。

配方验证：--deterministic-mode

应用规范环境设置，启用 PyTorch 确定性算法，并拒绝没有确定性路径的功能。内核测试：内核测试使用相同的输入和 RNG 状态重复操作，然后逐字节比较输出和梯度。模块验证：模块级测试在并行配置（包括 FP8 和 FP4）下，使用恢复的 RNG 状态重复模型和 transformer 块。随后，独立的端到端运行验证全精度训练指标是否保持逐位一致。

有关当前覆盖范围和已知缺口，请参阅确定性状态、操作目录和内核测试指南。

非确定性的常见来源

使用表 2 识别非确定性的可能原因，并选择下一步诊断方法。

观察到的症状 可能原因 诊断方法 解决方案
运行从第一步就出现分歧，但不一致 运行时自动调优选择了不同的内核配置 比较所选配置并追踪最早受影响的输出 固定或缓存已验证的配置
MoE 运行仅在启用融合辅助损失时出现分歧 归约顺序未固定 禁用单个融合并追踪损失计算 使用确定性归约或禁用融合
恢复的运行逐渐与连续运行分离 RNG 状态恢复不完整 在恢复前后立即比较 RNG 状态 逐位保存和恢复每个 RNG 流
软件或容器更新后出现分歧 底层库内核发生了变化 对堆栈进行二分查找并构建最小内核复现器 选择确定性内核路径，直到问题修复
加载检查点后参数立即不同 低精度权重或缩放因子以不同方式重建 比较保存/加载过程中的值和缩放元数据 保留全精度真实来源和精确重建路径
配方在小规模下通过但在大规模下失败 规模相关的通信或专家并行路径 系统地增加规模并追踪选定的 rank 隔离并修正第一个规模相关的操作

表 2. 常见的确定性症状、原因、诊断方法和解决方案

开始使用

逐位确定性使大规模预训练更易于调试和恢复。Nemotron 案例研究展示了内核和配方优化如何降低其性能成本。在您计划使用的硬件和软件环境中验证可复现性和开销。

要开始使用，请查看 Megatron Core 用户指南以获取受支持的配方，然后比较独立运行和检查点恢复运行以验证可复现性。