用NVIDIA Megatron Core实现大规模确定性预训练

位级确定性让大规模预训练更易于调试、验证并可复现地恢复。当训练参数规模达到数万亿、跨数千个GPU运行时，多维并行、低精度计算与分布式检查点会使故障复现和修复验证变得复杂，此时这些优势尤为重要。

生产级和大型训练任务成本高昂，需要高成功率保障。位级确定性能够支持可靠的故障重放、损失突变调试、系统变更验证以及中断训练的恢复。在万亿参数规模下，即便是微小的速度下降也可能耗费数千GPU天，因此确定性训练必须保持高效。一项万亿参数Nemotron案例研究展示了NVIDIA如何降低Megatron Core中位级确定性预训练的开销。

什么是位级确定性

位级确定性要求：当数据顺序、架构、训练方案、并行策略、软件、运行时设置和硬件均固定时，独立运行与从检查点恢复的运行遵循相同的数值轨迹。Megatron Core针对两种保证进行设计。

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

位级检查点恢复：保存并恢复一个或多个检查点的运行，必须与未中断的运行保持一致。

为什么预训练需要确定性

位级确定性为大规模预训练带来四项优势。

重放故障：在相同步骤复现损失突变，然后逐一更改组件以定位原因。

保持检查点连续性：恢复中断的训练而不改变数值轨迹。

检测异常差异：利用本应相同的运行之间出现的不匹配，排查检查点损坏或软硬件问题。

比较变更：区分预期更改带来的数值影响与运行间的正常波动。

如何验证确定性

打印的损失值不足以精确验证位级数值一致性。两次运行可能打印相同的四舍五入损失值，但梯度或参数的低位比特却存在差异。

更严格的验证流程是在每一步对选定的数值指标进行指纹记录，可包括训练损失、语言建模损失、负载均衡损失、多Token预测（MTP）损失以及梯度范数。将同一配置运行两次，逐步比较这些指纹。首个不匹配的步骤有助于缩小排查范围。

要验证位级检查点恢复，需将未中断的参考运行与从检查点恢复的运行进行比较。每次恢复后，验证参数、优化器状态、随机数生成器状态、数据位置以及后续输出是否与参考运行位级一致。

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

确定性被破坏时如何修复

在大规模场景下，非确定性可能是间歇性或依赖于拓扑结构的，可见的损失发散可能在首个差异比特出现很久之后才显现。Megatron-LM PR #7262中提出的工作流会记录按顺序排列的、按rank划分的张量指纹，用于离线比较。首先，需确认种子、数据顺序、批大小、并行策略、容器及软件栈完全一致。

通过细粒度追踪定位发散点

先从端到端比较开始，再将捕获范围从训练迭代逐步缩小至阶段、模块、操作、内核层级。如果两次运行以位级一致的输入进入某个范围，却以不同的输出离开，则发散源自该范围内部。

端到端指标用于检测故障，并确定序列化指标首次出现差异的迭代轮次。

广泛的语义追踪覆盖集合通信、流水线通信、重计算、优化器操作以及梯度归约，以确定发散发生的训练阶段。

模块与层级追踪将发散范围缩小至特定模型组件、Transformer层、混合专家（MoE）模块或优化器阶段。

操作级追踪对ATen操作及目标扩展调用进行指纹记录，以确定输入匹配但输出不同的首个操作。

内核级工具用于确定负责的内核、算法配置及非确定性机制。

不要只追踪损失明显发散的那一轮迭代；首个差异比特可能出现得更早。应先用广泛追踪定位最早出现发散的迭代轮次和rank，再在缩小的范围内启用详细的操作和内核追踪。

离线比较追踪结果

每个选定的rank将其追踪结果写入文件，不添加集合通信或跨rank排序，以降低掩盖所调查竞态条件的风险。按与运行无关的身份标识（如操作名称、出现次数及模块范围）对齐事件。

寻找输入指纹匹配但输出指纹不同的情况。

hash(输入_运行A) == hash(输入_运行B)，hash(输出_运行A) != hash(输出_运行B)

“首次不匹配”是针对每个rank分别判定的；序列号无法确立跨rank的顺序关系。根据因果角色对每个rank的首次不匹配进行分类。

如果多个rank指向同一个源头操作，则该操作本身很可能是非确定性的。如果只有部分rank如此，则需调查拓扑结构、rank分布、输入分布及归约顺序。

在GPU上进行张量指纹记录

该工作流使用torch.hash_tensor生成驻留于GPU的指纹。需记录每个张量的形状、数据类型及元素数量，并附上其摘要值。对于MXFP8或NVFP4张量，应同时对编码值及其缩放因子缓冲区进行指纹记录。

整张量的XOR指纹无法检测排列顺序的变化。对于路由映射、MoE分发输出等顺序敏感的数据，应使用dim参数按行或按块进行指纹记录。指纹只是一种高效的检验手段，匹配并不能保证张量完全位级一致。应使用字节级比较来确认位级相等性。

排除误报

在确定根本原因之前，需排查追踪盲区和无效读取情况。

调度器盲区：TorchDispatchMode只能观察到通过PyTorch调度器路由的ATen操作，自定义内核可能逃脱该追踪范围。如果首个不匹配出现在简单的视图、切片或加法操作上，应探查产生该输入的自定义内核。

探测伪影：未初始化的张量、未完成的异步集合通信以及非阻塞拷贝，可能在内容有效之前就被读取。应排除这些情况后，才能判定生产者为非确定性的。

验证修复效果

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

原始配对运行应能复现发散现象。

修复后的配对运行应保持位级一致。

判断该修复是纠正了非确定性实现，还是绕开了该实现。

测量新增的性能开销。

需注意，位级确定性验证仅在同一硬件与软件环境内有效。跨GPU代际、网络配置或库版本的比较超出本验证范围，可能产生不同结果。

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

一个开销巨大的确定性方案可以辅助调试，但对生产训练而言成本过高。确定性与非确定性执行应从一开始就协同优化。该研究使用的万亿参数Nemotron模型结合了Mamba风格的状态空间模型（SSM）层与Transformer注意力层。确定性训练必须覆盖SSM与注意力内核、低精度计算、分布式通信以及层类型之间的转换。

建立受控基线

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

测量指标包括：

吞吐量与单步耗时

峰值内存与GPU利用率

暴露的通信时间

主要内核组所耗费的时间

计算确定性开销：

确定性开销 = (确定性单步耗时 / 基线单步耗时 – 1) × 100%

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

优化历程

通过对三种工作负载进行优化，将确定性开销从两位数的基线水平降低至个位数。在2432个GPU的规模下，大规模Nemotron方案测得稳态确定性开销约为2%，并在800步训练中保持了位级确定性。

内核级优化示例

在分组GEMM的收尾阶段，原本多个N分块会累加写入同一个dprob[token]地址，导致结果依赖于写入到达的顺序。将这些写入操作串行化可以恢复确定性，但会降低并行度。

解决方案是为每个N分块分配独立的输出槽位，在内核内部保持并行执行，待所有写入操作完成后，再按固定顺序合并这些槽位的结果。

在大规模下验证优化效果

在少量GPU上测得的低确定性开销，并不能保证在生产规模下也能得到相同结果。通信、同步、流水线气泡、专家路由及负载均衡都可能改变相对开销比例。假设在10000个GPU上运行的非确定性基线需要100天，若将开销从15%降低到5%，则可节省10天，相当于节省10万GPU天。

随着模型、内核与方案演进维持确定性

新的内核、融合算子、精度格式或并行配置都可能引入非确定性。Megatron-LM通过以下检查机制防止出现回归问题。

方案验证：--deterministic-mode会应用标准环境设置，启用PyTorch确定性算法，并拒绝不具备确定性路径的功能。

内核测试：内核测试使用相同输入和随机数生成器状态重复执行操作，然后逐字节比较输出和梯度。

模块验证：模块级测试会在恢复随机数生成器状态的情况下，跨多种并行配置（包括FP8和FP4）重复运行模型及Transformer模块。随后通过独立的端到端运行，验证全精度训练指标是否保持位级一致。

关于当前覆盖范围及已知差距，可参阅确定性状态文档、操作目录及内核测试指南。

常见的非确定性来源

可参考相关对照表来识别可能导致非确定性的原因，并选择下一步诊断方法。

开始使用

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

如需开始使用，可查阅Megatron Core用户指南以获取支持的方案，然后比较独立运行与检查点恢复运行，以验证可复现性。

Q&A

Q1：位级确定性对大规模AI训练有什么用？

A：位级确定性能让大规模预训练更易于调试、验证和可复现地恢复，支持可靠的故障重放、损失突变调试、系统变更验证以及中断训练的恢复，这在训练万亿参数规模模型时尤为重要。

Q2：实现位级确定性训练会增加多少性能开销？

A：通过NVIDIA的优化，在2432个GPU规模的万亿参数Nemotron训练中，确定性训练带来的开销已降低到约2%的稳态水平，相比优化前两位数的开销有大幅改善。

Q3：如何验证两次训练运行是否真正达到位级一致？

A：不能仅依靠打印的损失值判断，因为四舍五入后的数值可能相同但底层比特不同。需要在每一步对训练损失、梯度范数等数值指标进行指纹记录，并逐步比较，找到首个不匹配的步骤来缩小排查范围。