用FP8訓練大模型有多香？ 微軟：比BF16快64%，省42%記憶體

原文來源：機器之心

圖片來源：由無界 AI生成

低精度訓練是大模型訓練中擴展模型大小，節約訓練成本的最關鍵技術之一。 相比於當前的 16 位和 32 位浮點混合精度訓練，使用 FP8 8 位浮點混合精度訓練能帶來 2 倍的速度提升，節省 50% - 75% 的顯存和 50% - 75% 的通信成本，而且英偉達最新一代卡皇 H100 自帶良好的 FP8 硬體支援。 但目前業界大模型訓練框架對 FP8 訓練的支援還非常有限。 最近，微軟提出了一種用於訓練 LLM 的 FP8 混合精度框架 FP8-LM，將 FP8 盡可能應用在大模型訓練的計算、存儲和通信中，使用 H100 訓練 GPT-175B 的速度比 BF16 快 64%，節省 42% 的記憶體佔用。 更重要的是：它開源了。

大型語言模型（LLM）具有前所未有的語言理解和生成能力，但是解鎖這些高級的能力需要巨大的模型規模和訓練計算量。 在這種背景下，尤其是當我們關注擴展至 OpenAI 提出的超級智慧 （Super Intelligence） 模型規模時，低精度訓練是其中最有效且最關鍵的技術之一，其優勢包括記憶體佔用小、訓練速度快，通信開銷低。 目前大多數訓練框架（如 Megatron-LM、MetaSeq 和 Colossal-AI）訓練 LLM 預設使用 FP32 全精度或者 FP16/BF16 混合精度。

但這仍然沒有推至極限：隨著英偉達 H100 GPU 的發佈，FP8 正在成為下一代低精度表徵的數據類型。 理論上，相比於當前的 FP16/BF16 浮點混合精度訓練，FP8 能帶來 2 倍的速度提升，節省 50% - 75% 的記憶體成本和 50% - 75% 的通信成本。

儘管如此，目前對 FP8 訓練的支援還很有限。 英偉達的 Transformer Engine （TE），只將 FP8 用於 GEMM 計算，其所帶來的端到端加速、記憶體和通信成本節省優勢就非常有限了。

但現在微軟開源的 FP8-LM FP8 混合精度框架極大地解決了這個問題：FP8-LM 框架經過高度優化，在訓練前向和後向傳遞中全程使用 FP8 格式，極大降低了系統的計算，顯存和通信開銷。

* 論文地址：

• 開源框架：

實驗結果表明，在 H100 GPU 平臺上訓練 GPT-175B 模型時， FP8-LM 混合精度訓練框架不僅減少了 42% 的實際記憶體佔用，而且運行速度比廣泛採用的 BF16 框架（即 Megatron-LM）快 64%，比 Nvidia Transformer Engine 快 17%。 而且在預訓練和多個下游任務上，使用 FP8-LM 訓練框架可以得到目前標準的 BF16 混合精度框架相似結果的模型。

在給定計算資源情況下，使用 FP8-LM 框架能夠無痛提升可訓練的模型大小多達 2.5 倍。 有研發人員在推特上熱議：如果 GPT-5 使用 FP8 訓練，即使只使用同樣數量的 H100，模型大小也將會是 GPT-4 的 2.5 倍！

Huggingface 研發工程師調侃：「太酷啦，通過 FP8 大規模訓練技術，可以實現計算欺騙！」

FP8-LM 主要貢獻：

• 一個新的 FP8 混合精度訓練框架。 其能以一種附加方式逐漸解鎖 8 位的權重、梯度、優化器和分散式訓練，這很便於使用。 這個8位框架可以簡單直接地替代現有16/32位混合精度方法中相應部分，而無需對超參數和訓練方式做任何修改。 此外，微軟的這個團隊還發佈了一個 PyTorch 實現，讓用戶可通過少量代碼就實現 8 位低精度訓練。
• 一個使用 FP8 訓練的 GPT 式模型系列。 他們使用了新提出的 FP8 方案來執行 GPT 預訓練和微調（包括 SFT 和 RLHF），結果表明新方法在參數量從 70 億到 1750 億的各種大小的模型都頗具潛力。 他們讓常用的並行計算範式都有了 FP8 支援，包括張量、流水線和序列並行化，從而讓使用者可以使用 FP8 來訓練大型基礎模型。 他們也以開源方式發佈了首個基於 Megatron-LM 實現的 FP8 GPT 訓練代碼庫。

FP8-LM 實現

具體來說，對於使用 FP8 來簡化混合精度和分散式訓練的目標，他們設計了三個優化層級。 這三個層級能以一種漸進方式來逐漸整合8位的集體通信優化器和分散式並行訓練。 優化層級越高，就說明 LLM 訓練中使用的 FP8 就越多。

此外，對於大規模訓練（比如在數千台 GPU 上訓練 GPT-175B），該框架能提供 FP8 精度的低位數並行化，包括張量、訓練流程和訓練的並行化，這能鋪就通往下一代低精度並行訓練的道路。

張量並行化是將一個模型的各個層分散到多台設備上，從而將權重、梯度和啟動張量的分片放在不同的 GPU 上。

為了讓張量並行化支援 FP8，微軟這個團隊的做法是將分片的權重和啟動張量轉換成 FP8 格式，以便線性層計算，從而讓前向計算和後向梯度集體通信全都使用 FP8。

另一方面，序列並行化則是將輸入序列切分成多個數據塊，然後將子序列饋送到不同設備以節省啟動記憶體。

如圖 2 所示，在一個 Transformer 模型中的不同部分，序列並行化和張量並行化正在執行，以充分利用可用記憶體並提高訓練效率。

而對於 ZeRO（零冗餘優化器 / Zero Redundancy Optimizer），卻無法直接應用 FP8，因為其難以處理與 FP8 劃分有關的縮放因數。 因此針對每個張量的縮放因數應當沿著 FP8 的劃分方式分佈。

為了解決這個問題，研究者實現了一種新的 FP8 分配方案，其可將每個張量作為一個整體分散到多台設備上，而不是像 ZeRO 方法一樣將其切分成多個子張量。 該方法是以一種貪婪的方式來處理 FP8 張量的分配，如演算法 1 所示。

具體來說，該方法首先根據大小對模型狀態的張量排序，然後根據每個 GPU 的剩餘記憶體大小將張量分配到不同的 GPU。 這種分配遵循的原則是：剩餘記憶體更大的 GPU 更優先接收新分配的張量。 通過這種方式，可以平滑地沿張量分配張量縮放因數，同時還能降低通信和計算複雜度。 圖 3 展示了使用和不使用縮放因數時，ZeRO 張量劃分方式之間的差異。

使用 FP8 訓練 LLM 並不容易。 其中涉及到很多挑戰性問題，比如數據下溢或溢出; 另外還有源自窄動態範圍的量化錯誤和 FP8 數據格式固有的精度下降問題。 這些難題會導致訓練過程中出現數值不穩定問題和不可逆的分歧問題。 為了解決這些問題，微軟提出了兩種技術：精度解耦（precision decoupling）和自動縮放（automatic scaling），以防止關鍵信息丟失。

精度解耦

精度解耦涉及到解耦數據精度對權重、梯度、優化器狀態等參數的影響，並將經過約簡的精度分配給對精度不敏感的元件。

針對精度解耦，該團隊表示他們發現了一個指導原則：梯度統計可以使用較低的精度，而主權重必需高精度。

更具體而言，一階梯度矩可以容忍較高的量化誤差，可以配備低精度的 FP8，而二階矩則需要更高的精度。 這是因為在使用 Adam 時，在模型更新期間，梯度的方向比其幅度更重要。 具有張量縮放能力的 FP8 可以有效地將一階矩的分佈保留成高精度張量，儘管它也會導致精度出現一定程度的下降。 由於梯度值通常很小，所以為二階梯度矩計算梯度的平方可能導致數據下溢問題。 因此，為了保留數值準確度，有必要分配更高的 16 位精度。

另一方面，他們還發現使用高精度來保存主權重也很關鍵。 其根本原因是在訓練過程中，權重更新有時候會變得非常大或非常小，對於主權重而言，更高的精度有助於防止權重更新時丟失信息，實現更穩定和更準確的訓練。

在該實現中，主權重有兩個可行選項：要麼使用 FP32 全精度，要麼使用帶張量縮放的 FP16。 帶張量縮放的 FP16 的優勢是能在無損於準確度的前提下節省記憶體。 因此，新框架的預設選擇是使用帶張量縮放的 FP16 來儲存優化器中的主權重。 在訓練中，對於 FP8 混合精度優化器，每個參數需要 6 個字節的記憶體：

相比於之前的解決方案，這種新的低位數優化器可將記憶體足跡降低 2.6 倍。 值得說明的是：這是首個用於 LLM 訓練的 FP8 優化器。 實驗表明 FP8 優化器能在從 1.25 億到 1750 億參數的各種模型大小下保持模型準確度。

自動縮放

自動縮放是為了將梯度值保存到 FP8 數據格式的表徵範圍內，這需要動態調整張量縮放因數，由此可以減少 all-reduce 通信過程中出現的數據下溢和溢出問題。

具體來說，研究者引入了一個自動縮放因數 μ，其可以在訓練過程中根據情況變化。

實驗結果

為了驗證新提出的 FP8 低精度框架，研究者實驗了用它來訓練 GPT 式的模型，其中包括預訓練和監督式微調（SFT）。 實驗在 Azure 雲計算最新 NDv5 H100 超算平台上進行。

實驗結果表明新提出的 FP8 方法是有效的：相比於之前廣泛使用 BF16 混合精度訓練方法，新方法優勢明顯，包括真實記憶體用量下降了 27%-42%（比如對於 GPT-7B 模型下降了 27%，對於 GPT-175B 模型則下降了 42%）; 權重梯度通信開銷更是下降了 63%-65%。

不修改學習率和權重衰減等任何超參數，不管是預訓練任務還是下游任務，使用 FP8 訓練的模型與使用 BF16 高精度訓練的模型的表現相當。 值得注意的是，在 GPT-175B 模型的訓練期間，相比於 TE 方法，在 H100 GPU 平臺上，新提出的 FP8 混合精度框架可將訓練時間減少 17%，同時記憶體佔用少 21%。 更重要的是，隨著模型規模繼續擴展，通過使用低精度的 FP8 還能進一步降低成本，如圖 1 所示。

對於微調，他們使用了 FP8 混合精度來進行指令微調，並使用了使用人類反饋的強化學習（RLHF）來更好地將預訓練后的 LLM 與終端任務和使用者偏好對齊。

結果發現，在Alpaca和MT-Bench基準上，使用FP8混合精度微調的模型與使用半精度 BF16 微調的模型的性能相當，而使用 FP8 的訓練速度還快 27%。 此外，FP8 混合精度在 RLHF 方面也展現出了巨大的潛力，該過程需要在訓練期間載入多個模型。 通過在訓練中使用 FP8，流行的 RLHF 框架 AlpacaFarm 可將模型權重減少 46%，將優化器狀態的記憶體消耗減少 62%。 這能進一步展現新提出的 FP8 低精度訓練框架的多功能性和適應性。

他們也進行了消融實驗，驗證了各元件的有效性。

可預見，FP8 低精度訓練將成為未來大模型研發的新基建。