在本博文中,我們將分解主要的 FP8 縮放策略 (按張量縮放、延遲和電流縮放以及按塊縮放 (包括 Blackwell 支持的 MXFP8 格式)) ,并解釋為什么每個策略都對在低精度訓練期間保持數值穩定性和準確性至關重要。了解這些方法有助于為您自己的 FP8 工作流程選擇合適的方法。
本文將探討 FP8 訓練的實際情況,重點介紹 NVIDIA Nemotron 實驗。我們將探討為什么需要按張量和延遲擴展,以及它們在哪些方面有不足,以及高級方法 (例如按塊子通道擴展) 如何為大型模型實現穩定、高效的 FP8 訓練。
在上一篇博文中,我們介紹了 FP8 精度,并將其兩種主要格式 E4M3 和 E5M2 與更成熟的類型(如 BF16)進行了比較。
按張量縮放
每張量縮放是一種基本的 FP8 策略,可為每個張量分配唯一的縮放系數(例如權重、激活函數或梯度),而不是使用單一的全局縮放。這一點至關重要,因為 FP8 的窄動態范圍無法處理跨不同張量的不同值分布。通過根據每個張量的統計數據調整縮放系數,每張量縮放有助于防止數值不穩定,并確保更準確、更穩定的訓練。
延遲縮放和當前縮放是每個張量縮放中的兩種常見方法。每種方法都提供了不同的縮放因子設置策略,但都在單個張量級別運行,以滿足 FP8 的獨特要求。
按張量延遲擴展
延遲縮放是一種實用的 FP8 縮放策略,該策略基于在先前迭代窗口中觀察到的最大絕對值 (amax) 來計算每個張量的縮放系數,而不僅僅是當前批量。這種基于歷史記錄的方法可以消除異常值,并降低可能會破壞訓練穩定的突然縮放變化的風險。
延遲擴展并非基于當前的數據批量來確定擴展系數,而是著眼于最近幾批中的最大值 (稱為“amax”) ,并保留這些最大值的運行歷史。
當選擇新的縮放系數時,系統會查看此歷史記錄并選擇反映近期趨勢的值,而不僅僅是最近的批量。這種方法有助于消除數據值的突然峰值或下降,使訓練過程更加穩定,并降低可能損害模型性能的突然變化的風險。

雖然延遲擴展適用于 FP8 訓練,但其有效性取決于一個關鍵假設:張量值的統計分布在訓練迭代中保持相對穩定。但是,異常值可能會在 amax 歷史記錄中占據主導地位,導致次優擴展,導致不足或溢出,并可能破壞訓練穩定,特別是在大規模運行中。
按張量電流縮放
雖然延遲擴展使用歷史數據,但每張量當前擴展提供了一種優先考慮實時適應性的替代方案。此方法根據當前向前或向后傳遞中各tensor的統計屬性確定其縮放系數。
按張量電流縮放可根據當前數據范圍動態調整縮放。這種即時響應能力有助于優化 FP8 表征,并可在訓練期間改善模型收斂。
其工作原理是在每個訓練批次中測量每個張量中的最大絕對值 (amax),無論是激活值、權重還是梯度。此測量是實時進行的,無需依賴前一批的任何歷史記錄。
然后,將縮放系數設置為 amax 適合 FP8 范圍,其中每個張量使用 FP8 的完整可用精度來執行特定步驟。
借助此縮放系數,張量的浮點值可量化為 FP8,從而在每次迭代中實現高效且準確的訓練。當前縮放的優勢如下:
- 瞬時適應:當前擴展可適應數據分布、學習率時間表或模型架構中的實時變化。這種實時響應能力可確保在擴展過程中不會出現延遲或不匹配。
- 簡單高效:缺少復雜的歷史緩沖區或跟蹤機制可簡化實施,并減少計算和內存用度。資源仍然直接集中在模型訓練上。
- 魯棒性:實時調整可針對當前批量中的異常值提高電流擴展的魯棒性。由于規模不受過去、過時或錯誤的 amax 值的影響,因此對于要求苛刻的深度學習工作流而言,它將提供更穩定的訓練體驗。

圖 2 突出顯示了在平均推理等具有挑戰性的基準測試中,與 BF16 基準相比,低精度格式的效率。
雖然 BF16 在整個 token 范圍內的總精度方面保持略微邊緣,但當前 scaling 緊隨其后,差距仍保持在較小的百分比余量內。
逐塊擴展

雖然每張量擴展方法提供了 FP8 訓練的基礎方法,但在處理單個張量內的塊級可變性時,它們通常會面臨局限性。
當張量內的不同區域或切片表現出不同的數值范圍和統計行為時,就會出現這種挑戰。例如,在大型 Transformer 模型中,單個權重矩陣中不同的注意力頭可能具有不同的數量級。同樣,高級層 (如 SwiGLU 或 mixture-of-experts (MoE) 路由器) 的激活函數因生成由小值和罕見但極端的異常值組成的密集集群的異構混合而聲名狼。
當對不同張量應用單個縮放系數時,要么犧牲密集、低數量級區域的精度,要么極端異常值可能會被削減。兩者都會降低訓練穩定性和模型的最終質量。
每塊縮放是一種功能強大的解決方案,可將每個 tensor 劃分為更小、更易于管理的連續塊,并為每個 張量分配專用的縮放因素。
縮放機制可以適應每個塊的局部統計數據,而不是由單個最極端值決定。這意味著高數量級區域得到準確表示,而不會影響同一張量內更小、更典型的值的保真度。
隨著深度學習模型的規模和架構復雜性不斷增加,細粒度、自適應擴展對于保持數字穩定性以及釋放 FP8 在現代硬件上提供的效率和速度至關重要。
接下來,我們將探索如何通過硬件原生 MXFP8 和其他可配置塊擴展方法等領先方案,在實踐中實施這一先進的擴展概念。
什么是 Micro-Scaling FP8?
Micro-Scaling FP8 (MXFP8) 基于逐塊擴展概念構建,代表了 NVIDIA Blackwell 硬件級解決方案,可實現高效穩定的 FP8 訓練。
MXFP8 符合更廣泛的 MX 數據格式標準,該標準定義了一個通用框架,用于跨各種低精度格式 (FP8、FP6、FP4 和 Int8) 進行共享的細粒度塊擴展。該標準具有三個關鍵組成部分:
- 縮放 (X) 數據類型/編碼:定義縮放系數本身的表示方式。
- 元素 (P_i) 數據類型/編碼:指定每個塊內實際數值元素的數據類型。
- 縮放塊大小 (k) :確定共享通用縮放系數的標量元素的數量。

MXFP8 的工作原理是什么?
MXFP8 實現塊級擴展,具有針對 Blackwell 架構優化的特定特征:
- 塊級除法:張量被系統劃分為連續塊,每個塊包含 32 個連續值,與 NVIDIA Blackwell Tensor Cores 的設計相匹配。
- 僅指數級縮放系數:32 值塊中的每個塊都分配了一個以 E8M0 格式存儲的專用縮放系數(8 位指數,0 位尾數),并表示一個 2 的冪乘數。這可優化硬件實現,并保持有利于深度學習的數值屬性。
- 轉置的硬件重構:鑒于縮放發生在特定方向 (例如,逐行或逐列) ,MXFP8 張量及其轉置在數字上不等效。當運算需要張量及其轉置 (常見于向前和向后傳遞) 時,Blackwell 單元會自動從高精度輸入進行重量化。這可避免因轉置量化 FP8 張量而復合量化錯誤,從而保持準確性。
這種創新的塊級擴展使 MXFP8 能夠克服按張量擴展的限制,適應單個張量內的數量級變化。通過保留高和低數量級區域,MXFP8 可最大限度地利用 FP8 的動態范圍并最小化量化誤差。這轉化為更快的訓練速度和更高的準確性,對于大型模型來說尤其如此。

圖 5 比較了 Nemotron 2b 和 8b 模型的 BF16 和 MXFP8 損失。在 2B 參數 Nemotron 模型中,MXFP8 顯示的驗證損失曲線與 BF16 相同,即使訓練令牌數量增加也是如此。這意味著,對于此模型,使用 BF16 和 MXFP8 數據類型之間沒有顯著的準確性差異,并且模型會按預期收斂。
借助 8B 參數模型,即使訓練令牌數量增加,MXFP8 的驗證損失曲線也與 BF16 相同。這再次表明,MXFP8 的收斂與 BF16 相似,即使模型大小增加,精度也沒有任何重大變化。
塊縮放
除了按張量縮放策略和 MXFP8 等特定于硬件的塊縮放之外,通用 FP8 塊縮放是一種通用且可配置的細粒度精度控制方法。
這項技術因其適應廣泛的模型架構和硬件要求而脫穎而出。

用戶可以指定各種形狀 (例如 1×128 或 128×128) ,以優化不同張量結構的縮放。這些特定于塊的擴展因子通常以 FP32 格式存儲,從而提供可靠的數值屬性。
這種可配置粒度可更精確地適應不同模型組件的獨特統計行為,成為跨不同工作負載優化 FP8 訓練的強大工具。
塊擴展的工作原理是什么?
FP8 塊縮放背后的核心原則是通過調整縮放系數以適應局部數據特征來提高精度。與在整個張量上應用單一比例的每張量方法不同,塊縮放會將每個張量分為更小、不同的片段。在每個定義的塊內,所有值都共享一個單獨存儲在 FP32 中的通用縮放系數,以確保準確性。
量化按元素塊的系數來擴展元素,并將其轉換為 FP8;去量化通過將 FP8 值乘以相同的系數來重建原始值。塊大小的選擇至關重要:較小的塊通過更精細的粒度減少量化誤差,而較大的塊則以增加噪音為代價降低存儲開銷。
NVIDIA NeMo 框架中的 Recipes
Transformer 引擎 FP8 配置通過高級配置公開,例如在 NVIDIA NeMo 框架中。用戶可以選擇不同的 fp8_recipe
標志來進行縮放策略,通常結合 BF16 進行混合精度訓練。
- 延遲擴展:
fp8_recipe = "delayed"
- 每張量電流縮放:
fp8_recipe = "tensorwise"
MXFP8:fp8_recipe = "mxfp8"
通用塊縮放:fp8_recipe = "blockwise"
總結
FP8 訓練是適用于大規模深度學習的實用高效解決方案。充分發揮 FP8 潛力的關鍵在于選擇和實施擴展策略。通過訪問 FP8 recipes 探索這些技術,并開始使用實用的 FP8 訓練配置和代碼。
?
?