前沿大语言模型的预训练本质上是吞吐量的竞赛。当训练规模跨越数万亿 token、动用数千个加速器时,每一步时间的微小百分比差异都可能累积成数天的额外训练耗时和巨大的算力开销。数值精度是影响训练效率的关键杠杆之一,但低位混合精度预训练在工程上极难调优。NVIDIA 近期在开发者博客上公开了基于 Blackwell 架构的 NVFP4 训练方案,为这一问题提供了新的解法。

该方案的核心是 TransformerEngine 中集成的 NVFP4 训练配方,它使 JAX 框架下的预训练能够使用亚字节精度。MaxText 作为高性能、可扩展的大模型库,提供了端到端的实现示例。结果显示,在 NVIDIA Blackwell 上进行的 4 位混合精度预训练,与 FP8 基线相比,吞吐量大幅提升,同时未出现可测量的精度损失。

NVFP4 格式本身采用了两级微缩放编码,相比其他微缩放格式,能以更低的误差传递更强的信号。在硬件层面,NVIDIA GB300 Grace Blackwell Ultra Superchip 对 NVFP4 的原生支持,使其 GEMM 吞吐量达到 Hopper 架构原生 FP8 精度的 7 倍。这种吞吐量跃升,配合精心设计的训练配方,直接缩短了每一步的训练时间,让 AI 工厂在同等时间预算下能够训练更多、更大的模型,或在更短的时间内完成模型训练。

训练配方的设计体现了对精度与速度的精细权衡。NVFP4 量化仅应用于 Transformer 模型中的 MLP(前馈网络)层,包括前向传播、激活梯度和权重梯度三个 GEMM 操作。注意力模块中的 QKV 投影、注意力输出投影以及分数与上下文的矩阵乘法均保持较高精度。这一选择基于一个关键观察:注意力机制中的 softmax 函数会指数级放大 QK^T 分数上的量化噪声,且注意力激活值中常含有集中分布的异常值,4 位精度难以有效表达。由于 MLP 层占据了训练过程中绝大部分浮点运算量,将 NVFP4 集中应用于此,既能捕获主要的加速收益,又规避了收敛风险。

在 MLP 层的 GEMM 内部,数据流经过精心编排。所有三个 GEMM 接收 NVFP4 输入,输出 BF16 格式的结果,最终在优化器步骤中汇入 FP32 主权重。配方中还引入了多项收敛保持技术:对权重采用二维块量化,以确保前向传播和激活梯度计算在转置操作下数值一致;在权重梯度输入上施加随机哈达玛变换,在 4 位量化前压平异常值;对梯度量化器使用随机舍入,使微小更新保持无偏。这些技术协同工作,使得 4 位预训练在维持所需精度的同时,充分释放了 Blackwell 与未来 Rubin 平台的 NVFP4 吞吐潜力。

实际性能数据验证了方案的有效性。基准测试使用 MaxText 对 Llama 3 8B 和 Llama 3.1 405B 进行预训练,序列长度设为 8192,每设备批次大小为 4,运行 50 步,并在相同的硬件、并行策略和全局批次大小下对比 NVFP4 配方与 FP8 基线。在 NVIDIA GB200 Grace Blackwell Superchip 和 GB300 Grace Blackwell Ultra Superchip 上,NVFP4 为每 GPU 带来了额外 500 至 700 TFLOPS 的持续计算能力。整体训练吞吐量较 FP8 基线提升 1.31 至 1.73 倍,其中 405B 模型的相对增益最为显著,在 GB200 上达到 1.44 倍,在 GB300 上达到 1.73 倍。这是因为 405B 模型每步的 GEMM 计算量远超 FSDP 集合通信开销,精度层面的加速直接转化为实际训练时间的节省。

精度方面,对 Llama 3 8B 进行 10000 步预训练的损失曲线显示,NVFP4 与 FP8 基线几乎完全重叠,从约 12.2 nats 下降至约 3.9 nats,收敛阶段的平均差距仅为 0.026 nats,完全处于步间噪声范围内。这意味着 NVFP4 带来的显著加速并未以任何可测量的精度损失为代价。

该方案已通过 JAX-Toolbox GitHub 仓库公开,开发者可在 Blackwell 硬件上拉取 MaxText 容器并运行示例脚本。脚本通过设置量化标志切换至 NVFP4 路径,并输出步时、每设备 TFLOPS 和每秒 token 数等关键性能指标,同时生成 Nsight Systems 跟踪文件以供深入分析。这一工具链的开放,为 AI 研究者和工程师在 Blackwell 平台上探索高效大模型训练提供了直接入口。

从产业角度看,NVFP4 的出现进一步拉高了 AI 训练基础设施的效率标杆。在算力成本居高不下、模型规模持续膨胀的背景下,能够在硬件层面原生支持超低精度训练,并在软件栈上提供开箱即用的收敛保障,意味着云服务商和 AI 实验室可以用更少的 GPU 小时完成同等规模的训练任务。对于 NVIDIA 自身而言,Blackwell 平台借此强化了其在 AI 芯片领域的代际优势,也为后续 Rubin 平台的精度演进铺平了道路。