跳到主要内容
返回时间线
NVIDIA Technical Blog来源发表:

万亿参数预训练也能逐位复现:NVIDIA把确定性开销从两位数压到约2%

核心概要

NVIDIA在Megatron Core上把逐位确定性预训练的开销压到可接受范围:同一初始状态下的两次独立运行、以及从检查点恢复的运行,都要在每一步保持数值完全一致。此前这类确定性配方开销常达两位数百分比,Nemotron 3 Ultra经缓冲区填充优化后从约17%降到2%,再降到1.5%(3,072块GPU);大规模Nemotron配方在2,432块GPU上测得约2%稳态开销,并在800步内保持逐位一致。

Source-provided article image: Scale Bitwise-Deterministic Pretraining with NVIDIA Megatron Core
Figure 1 · NVIDIA 开发者技术博客

深度剖析

万亿参数规模的预训练现在可以在保持逐位确定性的同时接近非确定性配方的速度,使失败重放、损失尖峰定位和中断续训都建立在可复现的数值轨迹上。 确定性配方过去常带有两位数百分比的开销,足以让生产级训练难以承受;而打印出的损失值精度不够,两次运行可能打印相同损失却在梯度和参数的低位比特上不同。 三个工作负载的优化把步时确定性开销从两位数基线降到低个位数:Nemotron 3 Ultra从约17%降到2%,再降到3,072块GPU上的1.5%;混合Triton代理从约60%降到37%,再降到2%;CuTeDSL权重梯度路径在256块GPU上约3.6%。

确定性通过固定归约顺序和保留全精度真值来实现,而不是靠事后比对损失曲线。 分组GEMM的收尾阶段原本让多个N-tile累加到同一个dprob[token]地址,结果依赖它们的到达顺序;串行化写入虽恢复确定性却损失并行度。 每个N-tile获得私有输出槽,核内并行保留,所有写入完成后按固定顺序合并;检查点路径保留全精度源真值,并用训练时相同的转换路径重建运行时表示。大规模Nemotron配方在2,432块GPU上测得约2%稳态开销,800步内保持逐位一致。

启示与展望

这套做法面向在数千块GPU上做万亿参数预训练、需要可复现失败重放和检查点续训的团队,验证在同一硬件与软件环境内完成。Megatron Core用户指南给出受支持的配方,随后可用独立运行与检查点恢复运行的对比来验证可复现性。

逐位确定性只在同一硬件和软件环境内得到验证,跨GPU代际、网络配置或库版本的比较不在验证范围内,可能产生不同结果。少数GPU上的低确定性开销不保证生产规模下相同,通信、同步、流水线气泡、专家路由和负载均衡都会改变相对开销;新内核、融合、精度格式或并行配置也可能重新引入非确定性。