MLSys’25 | 極低內存消耗：用SGD的內存成本實現AdamW的優化性能
本文第一作者 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 為資源受限的團隊提供了一條低成本、高效率的大模型訓練之路，預示著未來大模型訓練將進入一個全新的時代。
