©作者 |席浩誠
單位 | UC伯克利博士生
來源 | 機器之心
近期DeepSeek V3 引爆國內外的社交媒體，他們在訓練中成功應用了 FP8 精度，顯著降低了 GPU 記憶體使用和計算開銷。這表明，FP8 量化技術在最佳化大型模型訓練方面正發揮著越來越重要的作用。
近期，來自伯克利，英偉達，MIT 和清華的研究者們提出了視訊記憶體高效的 FP8 訓練方法：COAT（Compressing Optimizer states and Activation for Memory-Efficient FP8 Training），致力於透過 FP8 量化來壓縮最佳化器狀態和啟用值，從而提高記憶體利用率和訓練速度。
COAT 實現了端到端記憶體佔用減少 1.54 倍，端到端訓練速度提高 1.43 倍，同時保持模型精度。它還可以使訓練批次大小加倍，從而更好地利用 GPU 資源。
透過利用 FP8 精度，COAT 使大型模型的高效全引數訓練在更少的 GPU 上成為可能，並有助於在分散式訓練環境中加倍批次大小，為大規模模型訓練的擴充套件提供了實用的解決方案。最重要的是，他們的訓練程式碼完全開源。
論文第一作者席浩誠本科畢業於清華大學姚班，目前在伯克利攻讀博士學位，他在英偉達實習期間完成了這篇工作。論文共同通訊作者為 MIT 韓松副教授和清華大學陳鍵飛副教授。
論文題目：
COAT: Compressing Optimizer States and Activation for memory efficient FP8 Training
論文連結：
https://arxiv.org/abs/2410.19313
程式碼連結：
https://github.com/NVlabs/COAT
FP8最佳化器狀態
FP8 量化最佳化器狀態的難點
論文作者發現，當前的量化方法無法充分利用 FP8 的表示範圍，因此在使用每組量化（per-group quantization）對最佳化器狀態進行量化時會導致較大的量化誤差。
對於 FP8 的 E4M3 格式，我們希望量化組 X 的動態範圍覆蓋 E4M3 的最小可表示值（0.00195）和最大可表示值（448）之間的整個跨度，以充分利用其表示能力。
然而，E4M3 的動態範圍通常未被充分利用：E4M3 的動態範圍約為 200000，但一階動量的每個量化組的最大值最小值之比通常為 1000，二階動量的該比值則通常為 10，遠小於 E4M3 的動態範圍。這使得用 FP8 來量化最佳化器狀態的誤差非常大。
解決方案：動態範圍擴充套件
論文作者發現，在量化之前引入一個擴充套件函式 f (・)，能夠擴大量化組的動態範圍，並使其與 E4M3 對齊。使用的擴充套件函式為：
其中，k 是即時計算的引數，每個量化組共享一個 k。當 k > 1 時，動態範圍將被擴大，並更接近 E4M3 的動態範圍。在每一步訓練中，都可以即時的計算出最優的 k，從而可以充分利用 E4M3 的表示範圍，而原始的量化方法只能利用其中的一小部分。
動態範圍擴充套件方法可以大大減少量化誤差，並充分利用 E4M3 的動態範圍。除此之外，還發現，E4M3 比 E5M2 更適合一階動量。
而對於二階動量，雖然在原始設定中 E4M3 優於 E5M2，但在應用我們的擴充套件函式後，它們的量化誤差幾乎相同。因此，建議在量化最佳化器狀態時使用 E4M3 + E4M3 量化策略或 E4M3 + E5M2 量化策略。
FP8啟用
動機：非線性層佔用大量記憶體
在語言模型的前向傳播中，必須保留啟用值以用於反向傳播計算梯度。在 Llama 模型系列中，非線性層通常佔記憶體佔用的約 50%。相比之下，線性層的貢獻不到 25%。因此，最佳化線性和非線性層以減少啟用記憶體佔用至關重要。
解決方案：混合粒度 FP8 精度流
FP8 精度流要求所有線性和非線性層的輸入和輸出採用 FP8 格式。透過直接以 FP8 格式儲存輸入張量用於反向傳播，這消除了額外的量化操作需求，從而減少了相關開銷。
FP8 精度流自然地將非線性和線性層的記憶體佔用減少了 50%，因為它們只需要儲存 FP8 啟用值，而不是 BF16。為了進一步提高該方法的準確性，作者提出在不同層中變化量化粒度，以混合粒度的方式平衡精度和效率。
實驗結果
COAT 在多個任務中展示了其在記憶體佔用和訓練速度方面的優勢，同時保持了模型效能。
訓練加速1.43倍，視訊記憶體降低1.54倍
在使用 4 張 H100 訓練 Llama-2-13B 模型時，COAT 將每個 GPU 的記憶體佔用從 BF16 的 76.1GB 減少到 49.1GB，實現了 1.54 倍的記憶體縮減。
同時，COAT 將訓練速度從 BF16 的每秒 2345 個 token 提升至每秒 5295 個 token，達到 1.43 倍的加速。在幾乎所有的訓練場景下，COAT 都能夠使 Batch Size 翻倍，或是讓訓練所需的卡數減小。
訓練完全不掉點，FP8訓練表現和BF16吻合
COAT 在各種應用場景下，均展現出了出色的精度，完全不會導致模型效能下降。例如，在大語言模型預訓練任務中，COAT 可以保持近乎無損的模型效能，訓練中的 loss 曲線也和 BF16 完全吻合。
COAT 在視覺語言模型微調中同樣實現了和 BF16 訓練完全一致的表現。無論是 loss 曲線，還是下游任務上的表現，COAT 均和 BF16 基準相持平。
在一些實際的下游任務例子中，經過 COAT 訓練過的模型也有著相當優秀的生成和總結能力。
總結
COAT 的核心價值在於使用 FP8 進行訓練的同時做到了視訊記憶體最佳化。動態範圍擴充套件減少量化誤差，混合粒度量化最佳化啟用儲存，兩者協同作用使得端到端記憶體佔用降低 1.54 倍。
這種最佳化不僅適用於單機訓練，更在分散式訓練中發揮關鍵作用 —— 透過批次大小翻倍，可在相同硬體條件下處理更多資料，顯著提升訓練效率。而對於視訊記憶體資源緊張的研究者，COAT 也提供了全引數訓練的可行路徑，降低了大模型訓練的門檻。
更多閱讀
#投 稿 通 道#
讓你的文字被更多人看到
如何才能讓更多的優質內容以更短路徑到達讀者群體，縮短讀者尋找優質內容的成本呢？答案就是：你不認識的人。
總有一些你不認識的人，知道你想知道的東西。PaperWeekly 或許可以成為一座橋樑，促使不同背景、不同方向的學者和學術靈感相互碰撞，迸發出更多的可能性。
PaperWeekly 鼓勵高校實驗室或個人，在我們的平臺上分享各類優質內容，可以是最新論文解讀，也可以是學術熱點剖析、科研心得或競賽經驗講解等。我們的目的只有一個，讓知識真正流動起來。
📝 稿件基本要求：
• 文章確係個人原創作品，未曾在公開渠道發表，如為其他平臺已發表或待發表的文章，請明確標註
• 稿件建議以 markdown 格式撰寫，文中配圖以附件形式傳送，要求圖片清晰，無版權問題
• PaperWeekly 尊重原作者署名權，並將為每篇被採納的原創首發稿件，提供業內具有競爭力稿酬，具體依據文章閱讀量和文章質量階梯制結算
📬 投稿通道：
• 投稿郵箱：[email protected]
• 來稿請備註即時聯絡方式（微信），以便我們在稿件選用的第一時間聯絡作者
• 您也可以直接新增小編微信（pwbot02）快速投稿，備註：姓名-投稿
△長按新增PaperWeekly小編
🔍
現在，在「知乎」也能找到我們了
進入知乎首頁搜尋「PaperWeekly」
點選「關注」訂閱我們的專欄吧
·
·
·