前沿大語言模型的預訓練本質上是吞吐量的競賽。當訓練規模跨越數萬億 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 平臺的精度演進鋪平了道路。