MLSys’25 | 極低記憶體消耗：用SGD的記憶體成本實現AdamW的最佳化效能

机器之心發表於2025-02-27
原文網址 : https://www.jiqizhixin.com/articles/2025-02-27-4
記憶體
圖片
AIxiv專欄是機器之心釋出學術、技術內容的欄目。過去數年，機器之心AIxiv專欄接收報導了2000多篇內容，覆蓋全球各大高校與企業的頂級實驗室，有效促進了學術交流與傳播。如果您有優秀的工作想要分享，歡迎投稿或者聯絡報導。投稿郵箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com

本文第一作者 Hanqing Zhu 與 Zhengyu Zhang 均為 UT Austin 博士生。Hanqing Zhu 專注於高效 AI 計算，致力於最佳化機器學習硬體、系統及演算法；Zhengyu Zhang 聚焦構建高效可靠的機器學習系統。通訊作者為 UT Austin 的 David Z. Pan、Zhangyang Wang 以及 Meta AI 的 Jinwon Lee 。

大語言模型（LLM）正以驚人的速度重塑各個領域的格局，但其預訓練與微調過程卻面臨著前所未有的資源考驗 —— 不僅需要龐大的計算能力，更對記憶體提出了極高要求。

訓練 LLM 時，不僅需要儲存億級引數，還必須額外儲存梯度與最佳化器狀態（例如 Adam 中的動量和方差）。

例如，預訓練 LLaMA 7B 的一個批次就消耗高達 58GB 記憶體：14GB 用於模型引數，42GB 用於最佳化器狀態和梯度。這巨大的 “隱形” 記憶體開銷迫使研究者不得不選用視訊記憶體更大的高階 GPU、增加 GPU 數量，甚至犧牲訓練速度調低批次大小。

首次以類 SGD 記憶體成本完成大模型訓練

UT Austin 和 Meta AI 推出了全新訓練策略 ——APOLLO（Approximated Gradient Scaling for Memory Efficient LLM Optimization）。這一創新方法首次實現了以 SGD 級別的記憶體成本訓練大模型，同時保持了 AdamW 的卓越最佳化效能（已在 LLaMA 7B 的預訓練上驗證）。

文章已經被 MLSys 2025 接受，同時目前 APOLLO 已在 Hugging Face Transformers、LLaMA-Factory 等主流開源專案中落地，研究者可以輕鬆呼叫這一方法開展低記憶體大模型訓練與微調。

核心亮點

  • 極低記憶體消耗：首次以類 SGD 記憶體成本完成大模型訓練，達到甚至超越 AdamW 的效能。

  • 無需 SVD 計算：首次實現僅需輕量級隨機投影進行大模型預訓練，甚至在 7B 模型上最佳化速度超越 Adam。

卓越系統效能

  • 3 倍預訓練加速：在 8 塊 A100 GPU 上，APOLLO 預訓練 LLaMA 7B 模型實現了 3 倍的加速。

  • 突破規模限制：首次利用 DDP 成功訓練 13B 模型，並在 12GB 記憶體的消費級 GPU（如 NVIDIA RTX 4090）上完成 7B 模型的預訓練，無需依賴模型並行、檢查點或解除安裝策略。
圖片
圖片
  • 論文地址：https://arxiv.org/pdf/2412.05270
  • 論文網站：https://zhuhanqing.github.io/APOLLO/
  • 論文程式碼: https://github.com/zhuhanqing/APOLLO

方法簡要介紹

從 Adam 到結構化更新

在傳統的 Adam 最佳化器中，學習率的自適應更新是基於每個引數的梯度動量和方差：
圖片
APOLLO 首先將 Adam 的更新過程等效理解為對原始梯度 圖片 施加一個基於梯度動量和方差的梯度縮放因子圖片，進一步發現這一縮放因子可以透過結構化更新（例如按通道或張量進行縮放）來近似實現。

在 LLaMA-130M 上的實驗驗證表明，結構化更新不僅不會導致效能損失，反而會略微提升模型效能（見圖 1）。這一發現為 APOLLO 的設計提供了理論基礎。
圖片
圖 1: 結構化梯度縮放與 Adam 的訓練損失對比

APOLLO 框架

儘管結構化更新簡化了最佳化器設計，但仍需儲存完整的最佳化器狀態，記憶體開銷依然較高。這促使解決以下關鍵問題：

問題①：是否可以將結構化的學習率自適應轉換為一種既實用又記憶體高效的最佳化方法？

作者提出了 APOLLO，在低秩空間中構造一個輔助最佳化器狀態，用於預測結構化更新的梯度縮放因子（見圖 2）。由於最佳化器狀態建立在低秩空間中，其記憶體佔用顯著減少。
圖片 圖 2: APOLLO 框架

問題②：低秩空間中的自適應學習率能否有效近似其在原始空間中的行為？

APOLLO 透過理論證明，低秩空間中的結構化梯度縮放因子能夠有效近似其在原始空間中的行為。同時從低秩空間預測的縮放因子會以 √(n/r) 的比例變化（n 為原始空間維度，r 為低秩空間的秩）。

這使得 APOLLO 可透過放縮縮放因子來彌補低秩帶來的誤差， 從而可採用極低的秩近似，在不犧牲效能的前提下實現極低的記憶體消耗。實驗表明，在 LLaMA-7B 上，APOLLO 僅需 256 的秩，效能仍優於使用 1024 秩的 GaLore。
圖片
圖 3：梯度縮放因子隨秩的變化按照理論推導的 √(n/r) 變化

問題③：是否仍需要昂貴的 SVD 來構造低秩空間？

不需要。APOLLO 無需傳統方法中基於昂貴 SVD 運算的低秩投影方法，只需採用基於隨機投影的方法。這種方法避免了 SVD 的計算開銷，引入了極低的額外計算，甚至 APOLLO 在 7B 上的最佳化速度還快於 Adam。

問題④：是否能進一步壓縮最佳化器狀態至 SGD 的記憶體消耗，同時達到甚至超越 AdamW 的效能？

為實現極致記憶體高效，APOLLO 提出了 APOLLO-Mini。APOLLO-Mini 透過僅使用 rank-1 的輔助子空間進行張量級梯度縮放，而非 APOLLO 中的通道級更新。

實驗結果顯示，APOLLO-Mini 在實現 SGD 記憶體消耗的同時，其預訓練效能甚至超越了 AdamW。這一結果充分證明了結構化學習率更新的有效性，併為資源有限的研究者提供了高效訓練大模型的解決方案。

下圖為將 APOLLO 應用於 Adam 的演算法：
圖片
實驗及結果

研究者對 APOLLO 的 LLM 預訓練、微調以及系統級效能優勢進行了全面評估。

預訓練

研究團隊在 C4 資料集（經過淨化處理的 Common Crawl 語料庫）上對基於 LLaMA 的大型語言模型進行了預訓練實驗。

實驗結果顯示 APOLLO 系列方法（即使在使用極低秩的情況下，如秩分別為 256 和 1）均能顯著降低記憶體需求，並且獲得比 Adam 以及當前 SOTA 方法更好的結果。
圖片
表 1: 使用 C4 資料集預訓練不同大小的 LLaMA 模型驗證集困惑度（↓）對比

在 LLaMA 7B 模型上，APOLLO 系列不僅在記憶體使用上取得巨大突破，其訓練效能也遠超過 Adam 以及使用高秩（例如 1024）的 GaLore 方法。
圖片
表 2: 使用 C4 資料集預訓練 LLaMA-7B 模型驗證集困惑度（↓）對比

微調

微調任務（常識推理和 MMLU）中，APOLLO 的表現同樣搶眼，達到了甚至超越 Adam 基線的效果。
圖片
系統優勢

  • 加速訓練

在硬體資源受限的情況下，APOLLO 能夠支援比 AdamW 高達 4 倍的批次大小，從而實現最高可達 3 倍的訓練加速。

  • 極低記憶體消耗

首次實現了在 12GB 記憶體的消費級 GPU 上成功預訓練 LLaMA 7B 模型，為大模型訓練在低資源環境下提供了新的可能性。
圖片
圖 4：APOLLO 記憶體和訓練速度優勢

  • 極低的計算開銷

由於不依賴昂貴的 SVD 計算，APOLLO 最佳化時間在 7B 上甚至比 adam 還要快。

隨著大語言模型規模不斷攀升，記憶體需求也在呈爆炸式增長。APOLLO 不僅在理論上打破了最佳化器記憶體瓶頸，更在實踐中展現出預訓練與微調的卓越效能。APOLLO 為資源受限的團隊提供了一條低成本、高效率的大模型訓練之路，預示著未來大模型訓練將進入一個全新的時代。

相關文章