一、背景

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

對應(yīng)的論文：[2506.08027] Recipes for Pre-training LLMs with MXFP8 [1]

二、摘要

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

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

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

三、引言

3.1 MXFormat

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

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

• P：規(guī)定了 d 個 bit 數(shù)據(jù)的表示（編碼）方式，比如 FP8 的 E5M2 是怎么表示的。
• k：k 個元組作為一個 Block。
• X：上述 k 個元素的 Block 會對應(yīng)一個共享的 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 中也對上述問題進行了規(guī)范化，以 MXFP8 為例，其規(guī)定的 E4M3 和 E5M2 編碼方式如下圖 Table 1 和 Table 2 所示：

3.2 細(xì)粒度量化

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

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

更細(xì)粒度的量化也意味著更高的額外成本，Block 越?。６仍郊?xì)），額外成本越高。如下圖所示，對于一個常見的內(nèi)積操作：

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

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

如下圖所示，在 Tensor Core 計算時，可以將數(shù)據(jù) A/B 及它們對應(yīng)的 Scaling Factor A/B 一起輸入，并全部在 Tensor Core 內(nèi)完成。

當(dāng)然，其對數(shù)據(jù)類型也有一定的要求，如下圖所示：

四、MXFP8 預(yù)訓(xùn)練

4.1 轉(zhuǎn)換 FP32 到 MXFP8

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

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

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

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

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

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

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

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

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

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

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

4.2 所有 Tensor 采用 E4M3

在 Blackwell 架構(gòu)中，F(xiàn)P8 浮點格式包含兩種變體：E4M3 與 E5M2。實驗研究表明：

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

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

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

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

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

如下圖 Figure 4 所示，這種配置能最可靠地維持與 BF16 預(yù)訓(xùn)練相當(dāng)?shù)木人?，論文所有實驗均遵循此?zhǔn)則。在 MXFP 量化訓(xùn)練過程中，框架需為 Tensor（Weight、Activation 和 Gradient）保持兩個副本：每個副本沿點積歸約（dot-product reduction）軸（行與列）分別量化。Figure 4 展示了訓(xùn)練迭代中各 Tensor 在 Forward（FPROP）、Weight Gradient（WGRAD）和 Activation Gradient（DGRAD）計算中的使用方式。由于每個 Tensor 需以原始和轉(zhuǎn)置兩種形態(tài)參與運算，量化需沿行列兩個獨立軸向分別執(zhí)行。

當(dāng)前研究總結(jié)：提出的 MX Scale 因子舍入方案解決了基于 OCP 方法導(dǎo)致的不收斂問題，在 843M 參數(shù)模型上實現(xiàn)了 1T Token 訓(xùn)練下與 BF16 相當(dāng)?shù)木取＝Y(jié)合 Algorithm 1 中的 E4M3 格式及 Scale 因子計算方法，該方案可擴展至 8B 參數(shù)模型（15T Token 訓(xùn)練）——作者聲稱，這是目前采用 MXFP 格式的最大規(guī)模 LLM 預(yù)訓(xùn)練案例。

如下圖 Figure 6 所示，W16A2.5 規(guī)模 MoE 模型，1T Token 預(yù)訓(xùn)練也能實現(xiàn)同樣的效果：

4.3 15T Token MXFP8 預(yù)訓(xùn)練結(jié)果

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

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

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

如下圖 Figure 5 展示了 8B 預(yù)訓(xùn)練模型的訓(xùn)練損失及任務(wù)級準(zhǔn)確率?？梢钥闯觯瑑山M下游任務(wù)的評估分?jǐn)?shù)：

• MMLU 上的 5-shot 分?jǐn)?shù)。
• 9 個通用 Reasoning 基準(zhǔn)（ARC-Challenge 與ARC-Easy、Race、PIQA、Winogrande、Hellaswag、OpenBookQA、Social IQA 和 Commonsense QA）上 1-shot 分?jǐn)?shù)的平均值。

主要結(jié)果如下：

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

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

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

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

五、參考鏈接：

• [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

本文轉(zhuǎn)載自??AI閑談??，作者：AI閑談