NVFP4把4-bit预训练撑到10万亿token

Pretraining Large Language Models with NVFP4

NVIDIA, Felix Abecassis, Anjulie Agrusa, Dong Ahn, Jonah Alben, Stefania Alborghetti, Michael Andersch, Sivakumar Arayandi, Alexis Bjorlin, Aaron Blakeman, Evan Briones, Ian Buck, Bryan Catanzaro, Muya Chang, Jinhang Choi, Mike Chrzanowski, Eric Chung, Victor Cui, Steve Dai, Bita Darvish Rouhani, Carlo del Mundo, Deena Donia, Burc Eryilmaz, Henry Estela, Abhinav Goel, Oleg Goncharov, Yugi Guvvala, Robert Hesse, Russell Hewett, Herbert Hum, Ujval Kapasi, Brucek Khailany, Mikail Khona, Nick Knight, Alex Kondratenko, Ronny Krashinsky, Ben Lanir, Simon Layton, Michael Lightstone, Daniel Lo, Paulius Micikevicius, Asit Mishra, Tim Moon, Deepak Narayanan, Chao Ni, Abhijit Paithankar, Satish Pasumarthi, Ankit Patel, Mostofa Patwary, Ashwin Poojary, Gargi Prasad, Sweta Priyadarshi, Yigong Qin, Xiaowei Ren, Oleg Rybakov, Charbel Sakr, Sanjeev Satheesh, Stas Sergienko, Pasha Shamis, Kirthi Shankar, Nishant Sharma, Mohammad Shoeybi, Michael Siu, Misha Smelyanskiy, Darko Stosic, Dusan Stosic, Bor-Yiing Su, Frank Sun, Nima Tajbakhsh, Shelby Thomas, Przemek Tredak, Evgeny Tsykunov, Gandhi Vaithilingam, Aditya Vavre, Rangharajan Venkatesan, Roger Waleffe, Qiyu Wan, Hexin Wang, Mengdi Wang, Lizzie Wei, Hao Wu, Evan Wu, Keith Wyss, Ning Xu, Jinze Xue, Charlene Yang, Yujia Zhai, Ruoxi Zhang, Jingyang Zhu, Zhongbo Zhu

cs.CL, cs.AI, cs.LG

2025-09-30

NVFP4把12B模型训满10万亿token,稳定段相对FP8损失误差低于1%,MMLU-Pro 62.58%对62.62%。

这篇在解决什么

FP8 已经是大模型预训练的主流精度。再往下是 4-bit 浮点:Blackwell 上 FP4 Tensor Core 相对 FP8 的数学吞吐,GB200 约 2 倍、GB300 约 3 倍,操作数显存大约减半。代价是动态范围极窄,长 token 地平线上量化噪声会累积,训练容易飘、对不齐更高精度基线。

公开记录里,此前没有十亿参数级模型用 4-bit 真正跑过万亿 token 量级。这篇要回答的是算法问题,不是墙钟:NVFP4 加上一套针对性配方,能不能让 12B 混合 Mamba-Transformer 在 10 万亿 token 上跟上 FP8。

方法

NVFP4 相对 OCP 的 MXFP4 改了三处。块从 32 缩到 16,块级 scale 从整数幂 UE8M0 换成带尾数的 E4M3,再叠一层张量级 FP32 scale。块内最大绝对值可以按近 FP8 的精度还原,其余元素走 E2M1(可表示 0、±0.5 到 ±6 那一串离散点)。MXFP4 的幂次 scale 最坏会浪费 ±4、±6 两档,动态范围少将近一个 binade。

光换格式不够。所有线性层都打成 FP4 会发散。最终配方四条:

Embedding、输出头、归一化、非线性、attention 的 softmax 以及 QK、AV 批量 GEMM 都留在原精度。主权重、梯度累积、优化器状态走 FP32。Tensor 并行归约走 BF16。

结果

12B 模型沿用 Nemotron-H / Nemotron-Nano-12B-v2-Base 的混合架构(62 个 block:6 个 Self-Attention、28 个 FFN、28 个 Mamba-2),序列长度 8192,batch 736,WSD 学习率前 80% 恒定、后 20% 衰减,10 万亿 token。对照是同一套数据上的 FP8 预训练。下游评估一律在 BF16 做。

稳定段相对损失误差一直低于 1%,进入衰减后略超 1.5%。8T 处斜率变化来自学习率衰减,9T 处有一次数据配比切换。下游几乎贴着走:

任务FP8NVFP4
MMLU-Pro 5-shot62.6262.58
MMLU77.3676.57
GSM8k CoT89.0892.27
MATH83.3281.48
HumanEval+59.9357.43
MBPP+59.1155.91

编码三项落后约 2 到 3 点。论文怀疑 MBPP+ 是最后一个 checkpoint 的噪声,没有换 checkpoint 复核。想把损失缺口补回去,可以在衰减前切回更高精度:8.2T 之后切 BF16(约占训练 18%)能对齐 FP8;只在最后不到 1% 的步数切换,也能把相对误差从 1.5% 收到约 0.5%。

8B 模型、1 万亿 token 上对照 MXFP4:NVFP4 相对 BF16 的损失误差约 1.5%,MXFP4 约 2.5%。MXFP4 要多训 36% token(1.36T 对 1T)才能追上 NVFP4 的损失。

四件套拿掉任何一件,12B、10T 的收敛都会变差。小模型、短 horizon 上很多件看不出差别。

为什么重要

这是目前公开记录里最长的 4-bit 预训练。算法侧已经证明 NVFP4 能撑到 10T,工程侧 Transformer Engine 已经合入。在 Blackwell 上跑预训练的人,可以把 GEMM 密集的线性层下到 FP4,attention 和尾巴仍用更高精度。

它还不是全网 4-bit。16% 线性层留在 BF16,attention 路径没动。报告明确说关心的是算法可行性,不是端到端墙钟。真正省多少时间,取决于线性层占比和通信是否卡死。

相对 MXFP4,同样 token 预算下 NVFP4 收敛更好。少训 36% token 是能直接算的账。

局限与存疑

编码三项的落后没有得到清楚解释,只给了「最后 checkpoint 噪声」这一猜测。下游评估都在 BF16 做,看不到 FP4 推理时的落差。

配方对规模敏感:12B、10T 四件套都必要,1.2B 上很多件几乎没差。MoE、更长 horizon、把 attention 也量化,都还没做。最后几层必须留高精度,说明 FP4 对输出附近的动态范围仍然不够。

Hadamard 只打 Wgrad,是为了迁就二维权重缩放。更大模型的激活离群值会不会把 Fprop、Dgrad 也拖垮,这篇没验证。张量级 FP32 scale 会多一次全张量 amax 扫描,系统实现上不是免费的。

术语

原文与代码

社区讨论

相关论文

全部论文解读