一、背景

筆者之前寫過 FP8 訓練的綜述文章以及 FP4 訓練和推理的綜述文章，本文對其進一步補充，介紹 NVIDIA 最新的使用 MXFP8 預訓練的方案。

對應的論文：[2506.08027] Recipes for Pre-training LLMs with MXFP8 [1]

二、摘要

精度縮放——即在預訓練過程中使用更少的比特來表示模型參數及相關 Tensor——已成為一種在不犧牲精度前提下提升 GPU 效率的有效技術。NVIDIA 最新 Blackwell GPU 中引入 Microscaling (MX) 格式，為 Tensor 量化提供了細粒度解決方案。

盡管 MX 格式相較于其他低精度表示法有望提升數值穩定性，但在實際應用中仍需謹慎使用。本文研究表明：當采用 OCP 規范建議的舍入模式進行 LLM 預訓練時，會導致模型不收斂。為此，作者提出了一種改進的舍入模式——通過采用"向無窮大舍入"方式計算縮放因子，成功實現了 8B 參數模型在 15T Token 上采用 MXFP8 格式的預訓練。

PS：可能是因為采用 Hopper GPU 模擬的方式，而不是真實的 Blackwell 訓練，因此論文并沒有提供相應的效率提升數據。

三、引言

3.1 MXFormat

2023 年，OCP（Open Compute Project） 在 AMD, Arm, Intel, Meta, Microsoft, NVIDIA, Qualcomm 的參與下提出 Microscaling（MX）Format 規范（OCP Microscaling Formats (MX) Specification Version 1.0 [2]），主要是為了對跨硬件/軟件平臺可實施的新功能及格式進行標準化，有效減少軟件與基礎設施成本，并消除定制化解決方案帶來的各類附加費用或管理負擔，推動硬件性能與效率的提升。

如下圖所示，MX 最主要的特點是其包含三部分內容（很類似于常見的 Per-Block 細粒度量化方式，只不過這里是制定了一個統一的規范）：

• P：規定了 d 個 bit 數據的表示（編碼）方式，比如 FP8 的 E5M2 是怎么表示的。
• k：k 個元組作為一個 Block。
• X：上述 k 個元素的 Block 會對應一個共享的 Scale 值。

我們在之前的文章中提到過，即使都是 E5M2 或者 E4M3，不同公司的硬件可能采用不同的格式。比如 NVIDIA Hopper GPU 上的 E5M2 符合 IEEE 754 Style，而 E4M3 卻不符合 IEEE 754 Style。如下圖所示，IEEE 754 Style 的 E4M3 的范圍為 [-240, 240]，而 ARM-Intel-Nvidia Style 的 E4M3 的范圍是 [-448, 448]： 

在 MX 中也對上述問題進行了規范化，以 MXFP8 為例，其規定的 E4M3 和 E5M2 編碼方式如下圖 Table 1 和 Table 2 所示：

3.2 細粒度量化

更低的精度通常意味著更難量化，為了維持精度需要更細力度的 Scaling Factor，比如：

• FP16：早期使用 FP16 進行混合精度訓練時通常整個模型一個 Scaling Factor 即可。
• FP8：在 Inference 時通常 Per-Tensor 的 Scaling Factor 即可比較好的維持精度；而 Training 時往往需要 Per-Block 或 Per-Channel，不過 Block 通常比較大，比如 128x128 或 128x1。
• FP4：需要 Per-Block 量化，并且 Block 需要比較小，比如 32 或 16。

更細粒度的量化也意味著更高的額外成本，Block 越小（粒度越細），額外成本越高。如下圖所示，對于一個常見的內積操作：

• Per-Tensor 量化：需要額外執行 1 次 Scaling Factor 處理。
• Per-Block 量化：需要額外執行很多次 Scaling Factor 處理。

為了更好的解決上述問題，NVIDIA 在新的 Blackwell Tensor Core 中支持了新的 Block-Scaled 類型，原生支持 Microscaling Formats，如下圖 Table 1 所示，其支持 MXFP8、MXFP6、MXFP4：

如下圖所示，在 Tensor Core 計算時，可以將數據 A/B 及它們對應的 Scaling Factor A/B 一起輸入，并全部在 Tensor Core 內完成。

當然，其對數據類型也有一定的要求，如下圖所示：

四、MXFP8 預訓練

4.1 轉換 FP32 到 MXFP8

在訓練的 Forward 與 Backward 過程中，模型 Weight、Activation 及 Gradient Tensor 均由 FP32 量化到 MXFP8 格式。量化后的 MXFP8 Tensor 隨后存儲于硬件中并執行運算。作者首先闡述了轉換過程 Quantize_to_fp8(Vi/2X)。下文所述的量化方法統一適用于所有 MX 格式（包括 E4M3、E5M2、E2M3、E3M2 及 E2M1），僅 MX 數據類型存在差異。

計算 X 值：通常情況下，Tensor 中各 Block 內的大部分數值會超出目標 MX 格式的可表示范圍，既可能低于最小可表示數（下溢），也可能高于最大可表示數（上溢）。為解決這一問題，需將 Block 內所有數值乘以一個 Scale 因子，使絕大多數數值被調整至可表示范圍內。

該 Scale 因子 X 的計算基于 32 個（MX 規范，k）高精度輸入值中的絕對最大值（amax），即amax = max(‖Vi‖); 1≤i≤32。其核心目標是將輸入中的 amax 映射為 MX 格式中的最大可表示值。當輸入數據包含無窮大（Infinity）或非數值（NaN）時需特殊處理：若某 Block 的 amax 為 0，則設定 X=-127，此時 X 為 2-127，且該情況下所有 Qi’ 均置為 0。

根據 OCP 規范，當 X 不為 Inf、NaN 或 0 時，X 應設定為不超過 “amax 除以 MX 格式類型最大可表示二次冪” 的最大二次冪。以 E4M3 類型為例，由于 448 是其最大幅值，故 X 的計算公式為X=floor(log?(amax))/floor(log?(448))。值得注意的是，OCP 規范在此計算過程中忽略了該比率浮點數尾數部分的影響。

作者觀察到遵循 OCP 規范時存在精度下降現象。如下圖 Figure 2 所示為兩種 Token 規模（300B 和 1T）下訓練的 843M 參數量 Transformer 模型的訓練損失曲線。其采用兩種不同配置方案：

• cfg1：所有 Tensor（Weight W、Activation A、Gradient G）均采用 E4M3 格式。
• cfg2：Weight W、Activation A 采用 E4M3 格式，Gradient G 采用 E5M2 格式。

E5M2 格式相較于 E4M3 具有約 1.8 倍的 binades 優勢。鑒于 Gradient 通常具有更大的動態范圍，早期的工作 [2209.05433] FP8 Formats for Deep Learning [3] 主張采用 E5M2 格式進行 Tensor 縮放。實驗結果表明，在 cfg1 與 cfg2 中，使用 OCP 方法計算縮放因子均會導致訓練發散（如下圖 Figure 2a）或相對于 BF16 基準的損失差距擴大（如下圖 Figure 2b）。

如下圖 Algorithm 1 概述了作者計算 Scale 因子的方法。其核心改進在于：當處理 amax 與 MX 格式最大可表示值 destmax 的比值指數時，采用向正無窮方向的 round-up 策略（同時飽和至 UE8M0 格式的極值邊界）。這與 OCP 方案形成鮮明對比，后者實質上是建議對 Scale 值執行 round-down 操作。由于高精度值 Vi 需通過 Scale 因子 2X 進行縮放，對分數項 (Vi/2X) 分母實施 round-up 操作，會更傾向于將 amax 映射至 destmax 以下；反之，OCP 方法則傾向于使 amax 超過 destmax（后續必須通過截斷處理使其可表示）。作者推測 OCP 取整方法帶來的飽和效應會影響模型精度。

如上圖 Figure 2 所示，采用提出的舍入方案后，Gradient 位寬配置為 E4M3 的 MXFP8（藍色曲線）與 E5M2 的 MXFP8（紫色曲線）在 300B 和 1T Token 的訓練過程中，其損失曲線均與 BF16 完全重合。

FP32 數值到 MX 格式的量化過程：當縮放因子 X 確定后，Tensor Vi 通過乘以 2X 進行尺度變換，隨后量化至最接近的 FP8 可表示數值（即 Quantize_to_fp8()）。該量化步驟采用“就近取偶（Round-to-nearest-ties-to-even，RN）”舍入法，且轉換過程具有飽和特性——若舍入結果超出 FP8 最大值或低于最小值，則將結果截取至相應的極值。

這種轉換機制在低精度 LLM 預訓練中的典型應用場景是：矩陣乘積累加運算（MMA）的輸出（通常以 FP32 格式存儲）需要映射為 MXFP8 格式，相比存儲 FP32 數值可顯著節省寫入帶寬和存儲容量。模型后續運算讀取 MXFP8 數值時，相較加載 FP32 數據也能減少讀取帶寬消耗。此外，由于 Tensor Core 可直接處理 MX 格式輸入，低精度 MMA 操作不僅能降低能耗，還能獲得更高的計算吞吐量。

4.2 所有 Tensor 采用 E4M3

在 Blackwell 架構中，FP8 浮點格式包含兩種變體：E4M3 與 E5M2。實驗研究表明：

Weight 與 Activation 量化性能對比：采用 E4M3 格式量化 Weight 和 Activation 時展現出更優的訓練收斂性。如下圖 Figure 3a 所示（測試模型與 Figure 2b 相同，參數量 843M），當 Activation（紫色曲線）或 Weight（藍色曲線）采用 E5M2 格式時，其損失函數收斂性顯著差于所有 Tensor 采用 E4M3 量化方案（橙色曲線）。值得注意的是，僅 Gradient 采用 E5M2 時（黃色曲線）仍能維持較好的收斂特性。

Gradient Tensor 量化分析：E4M3 格式在 Gradient 量化中能保持與 BF16 預訓練相當的損失，這一優勢在參數量 ≥2B 的模型中尤為顯著。如下圖 Figure 3c 展示了 8B LLM（1T Token 訓練）的對比結果：E4M3 Gradient 量化（橙色曲線）的最終損失值顯著低于 E5M2 方案（黃色曲線），且該差距隨訓練 Token 數量增加而擴大。這一現象揭示了模型參數量對數值格式選擇的敏感性，強調需在不同規模模型中系統評估格式的數值特性（PS：這也是為什么筆者一直提到之前很多文章只在小規模模型、數據量下做實驗不夠有說服力的原因）。

既往研究采用的 Tensor 級縮放 FP8 方案及 DeepSeek V3（[2412.19437] DeepSeek-V3 Technical Report [4]） 提出的粗粒度 Block-Scaled 方案均默認選用 E5M2 格式處理 Gradient Tensor（PS：論文這里表述有問題，DeepSeek V3 中其實所有 Tensor 都已經采用 E4M3 格式；此外，這里的粗粒度是相對 32 的 Block 大小而言，在 DeepSeek 中為了效率采用的是 128 或者 128x128 的 Block 大小）。本研究發現：當采用細粒度縮放（32 元素 Block）時，E4M3 格式的 17.8 個 binades 可充分滿足動態范圍需求。在滿足動態范圍前提下，量化精度成為關鍵因素——E4M3 每個指數區間包含 8 個量化樣本，其采樣密度是 E5M2（4樣本/區間）的 2 倍。因此，提出的 MXFP8 預訓練方案對所有三類 Tensor（Weight、Activation、Gradient）均采用 E4M3 數據類型進行量化。

MXFP8 實例化層級及訓練流程：論文所有研究均采用基于語言的 Transformer 模型，未來工作將探索該方案在語音與視覺模型中的應用。研究表明，量化策略建議：

• 將模型中所有 Transformer  Block 的 QKV、Proj 以及 FFN 的 Up-proj 和 Down-proj 轉換為 MXFP8 格式。
• Self-Attention 中的批量矩陣乘法（BMM1：Query-Key 點積和 BMM2：Attention Score-Value 乘積）以及 Softmax、激活函數和殘差相加等運算仍保持高精度計算。
• 輸入 Embedding 層和最終輸出 LM-head 同樣采用 BF16 或 FP16 格式。

如下圖 Figure 4 所示，這種配置能最可靠地維持與 BF16 預訓練相當的精度水平，論文所有實驗均遵循此準則。在 MXFP 量化訓練過程中，框架需為 Tensor（Weight、Activation 和 Gradient）保持兩個副本：每個副本沿點積歸約（dot-product reduction）軸（行與列）分別量化。Figure 4 展示了訓練迭代中各 Tensor 在 Forward（FPROP）、Weight Gradient（WGRAD）和 Activation Gradient（DGRAD）計算中的使用方式。由于每個 Tensor 需以原始和轉置兩種形態參與運算，量化需沿行列兩個獨立軸向分別執行。

當前研究總結：提出的 MX Scale 因子舍入方案解決了基于 OCP 方法導致的不收斂問題，在 843M 參數模型上實現了 1T Token 訓練下與 BF16 相當的精度。結合 Algorithm 1 中的 E4M3 格式及 Scale 因子計算方法，該方案可擴展至 8B 參數模型（15T Token 訓練）——作者聲稱，這是目前采用 MXFP 格式的最大規模 LLM 預訓練案例。

如下圖 Figure 6 所示，W16A2.5 規模 MoE 模型，1T Token 預訓練也能實現同樣的效果：

4.3 15T Token MXFP8 預訓練結果

作者采用 Megatron-LM 框架預訓練了一個 8B 參數的 Nemotron 模型。該模型包含 32 個 Transformer Block，每個 Block 32 個 Attention Head，隱層維度為 4096，采用 GQA 且 Group 大小為 8，KV 通道數為 128，預訓練階段序列長度為 8192。共訓練 15T Token，Batch Size 為 768。初始學習率設為 6×10??，并通過 cosine decay 到 6×10??。如下圖 Table 2 為幾個模型的詳細配置：

采用分階段數據混合策略進行訓練：第一階段使用促進數據多樣性的混合數據集，第二階段則轉向高質量數據集（如維基百科），在訓練進度達到 60% 時切換至第二階段。此類混合策略在其他大規模預訓練框架中亦有應用。

模型預訓練在 3072 Hopper GPU 上完成（實驗周期內缺乏支持 MX 格式的 Bloackwell 硬件平臺）。通過在 Hopper GPU 上模擬 MX 格式實現：輸入矩陣乘法加速器（MMA）的 Tensor 先量化為 MX 格式，在執行 BF16 MMA 運算前轉換回 BF16 格式。為驗證模擬方案的數值保真度，作者與 Blackwell 平臺上采用真實 MXFP8 格式訓練的 2B 參數 LLM 進行對比實驗，確認二者輸出結果完全一致。

如下圖 Figure 5 展示了 8B 預訓練模型的訓練損失及任務級準確率。可以看出，兩組下游任務的評估分數：

• MMLU 上的 5-shot 分數。
• 9 個通用 Reasoning 基準（ARC-Challenge 與ARC-Easy、Race、PIQA、Winogrande、Hellaswag、OpenBookQA、Social IQA 和 Commonsense QA）上 1-shot 分數的平均值。

主要結果如下：

• 采用 MXFP8 預訓練時，模型的驗證困惑度與 BF16 預訓練結果持平（Figure 5 左圖）。在整個預訓練過程中，MXFP8 與 BF16 的驗證困惑度差異始終小于 0.50%。
• Figure 5 中、右兩圖顯示了兩組下游任務的評估分數。MXFP8 訓練模型的得分與 BF16 訓練模型完全匹配，證明 MXFP8 可作為 LLM 預訓練的有效候選方案。

MXFP8 與 FP8 對比：除 MXFP8 和 BF16 外，Figure 5 還展示了傳統 FP8 精度訓練同模型的任務級分數。FP8 方案采用軟件模擬的分塊縮放技術，通過整體 Tensor 縮放使多數 Tensor 值落入量化格式的可表示范圍。遵循 [2504.03624] Nemotron-H: A Family of Accurate and Efficient Hybrid Mamba-Transformer Models [5] 的 FP8 預訓練設置建議：模型首尾 Transformer Block 保持 BF16 精度，其余 Block 的線性層量化為 FP8，該配置適用于 20T Token 規模的 8B 和 56B 參數 LLM 預訓練。但保留部分 BF16 層會影響端到端加速比，并增加預訓練復雜性——需額外決策哪些層維持高精度。實驗表明，MXFP8 在這兩組任務上無需任何 BF16 層即可達到與 FP8 相當的精度。

MXFP8 與分塊 FP8 的對比：進一步地，諸如 Deepseek-V3 等研究表明，在使用 FP8 時需要縮小 Block 規模。在此配置下，部分 Tensor 需采用 1x128 向量級縮放，而其他 Tensor 則需實施分塊（如 128x128）縮放，這增加了 GEMM Kernel 函數設計的復雜度。MXFP8 的原生支持則簡化了這一過程——其細粒度縮放機制提供了更優的數值魯棒性，同時規避了小 Block 尺寸與硬件速度之間的權衡問題。

綜上所述，相比于 BF16 或 FP8 預訓練，MXFP8 能保持同等精度。在 GB200 Blackwell 系統上，MXFP8 的吞吐量是 BF16 的 2 倍，這使得端到端 MXFP8 預訓練速度超越 BF16 預訓練。與 FP8 相比，MXFP8 方案還更加簡便（所有層均可量化且縮放由硬件處理），同時保持同等或更優的吞吐性能。

五、參考鏈接：

• [1] https://arxiv.org/abs/2506.08027
• [2] https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf
• [3] https://arxiv.org/abs/2209.05433
• [4] https://arxiv.org/abs/2412.19437
• [5] https://arxiv.org/abs/2504.03624

本文轉載自??AI閑談??，作者：AI閑談