用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用户指南以获取支持的方案然后比较独立运行与检查点恢复运行以验证可复现性。QAQ1位级确定性对大规模AI训练有什么用A位级确定性能让大规模预训练更易于调试、验证和可复现地恢复支持可靠的故障重放、损失突变调试、系统变更验证以及中断训练的恢复这在训练万亿参数规模模型时尤为重要。Q2实现位级确定性训练会增加多少性能开销A通过NVIDIA的优化在2432个GPU规模的万亿参数Nemotron训练中确定性训练带来的开销已降低到约2%的稳态水平相比优化前两位数的开销有大幅改善。Q3如何验证两次训练运行是否真正达到位级一致A不能仅依靠打印的损失值判断因为四舍五入后的数值可能相同但底层比特不同。需要在每一步对训练损失、梯度范数等数值指标进行指纹记录并逐步比较找到首个不匹配的步骤来缩小排查范围。