用NVIDIA Megatron Core实现大规模确定性预训练|core|gpu|nvidia|确定性预训练|算法

位级确定性让大规模预训练更易于调试、验证并可复现地恢复。当训练参数规模达到数万亿、跨数千个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步训练中保持了位级确定性。