本文第一作者 Hanqing Zhu 與 Zhengyu Zhang 均為 UT Austin 博士生。Hanqing Zhu 專注于高效 AI 計算，致力于優(yōu)化機器學(xué)習(xí)硬件、系統(tǒng)及算法；Zhengyu Zhang 聚焦構(gòu)建高效可靠的機器學(xué)習(xí)系統(tǒng)。通訊作者為 UT Austin 的 David Z. Pan、Zhangyang Wang 以及 Meta AI 的 Jinwon Lee 。

大語言模型（LLM）正以驚人的速度重塑各個領(lǐng)域的格局，但其預(yù)訓(xùn)練與微調(diào)過程卻面臨著前所未有的資源考驗 —— 不僅需要龐大的計算能力，更對內(nèi)存提出了極高要求。

訓(xùn)練 LLM 時，不僅需要存儲億級參數(shù)，還必須額外保存梯度與優(yōu)化器狀態(tài)（例如 Adam 中的動量和方差）。

例如，預(yù)訓(xùn)練 LLaMA 7B 的一個批次就消耗高達 58GB 內(nèi)存：14GB 用于模型參數(shù)，42GB 用于優(yōu)化器狀態(tài)和梯度。這巨大的 “隱形” 內(nèi)存開銷迫使研究者不得不選用顯存更大的高端 GPU、增加 GPU 數(shù)量，甚至犧牲訓(xùn)練速度調(diào)低批量大小。

首次以類 SGD 內(nèi)存成本完成大模型訓(xùn)練

UT Austin 和 Meta AI 推出了全新訓(xùn)練策略 ——APOLLO（Approximated Gradient Scaling for Memory Efficient LLM Optimization）。這一創(chuàng)新方法首次實現(xiàn)了以 SGD 級別的內(nèi)存成本訓(xùn)練大模型，同時保持了 AdamW 的卓越優(yōu)化性能（已在 LLaMA 7B 的預(yù)訓(xùn)練上驗證）。

文章已經(jīng)被 MLSys 2025 接受，同時目前 APOLLO 已在 Hugging Face Transformers、LLaMA-Factory 等主流開源項目中落地，研究者可以輕松調(diào)用這一方法開展低內(nèi)存大模型訓(xùn)練與微調(diào)。

核心亮點

• 極低內(nèi)存消耗：首次以類 SGD 內(nèi)存成本完成大模型訓(xùn)練，達到甚至超越 AdamW 的性能。
  
  
• 無需 SVD 計算：首次實現(xiàn)僅需輕量級隨機投影進行大模型預(yù)訓(xùn)練，甚至在 7B 模型上優(yōu)化速度超越 Adam。

卓越系統(tǒng)性能

• 3 倍預(yù)訓(xùn)練加速：在 8 塊 A100 GPU 上，APOLLO 預(yù)訓(xùn)練 LLaMA 7B 模型實現(xiàn)了 3 倍的加速。
  
  
• 突破規(guī)模限制：首次利用 DDP 成功訓(xùn)練 13B 模型，并在 12GB 內(nèi)存的消費級 GPU（如 NVIDIA RTX 4090）上完成 7B 模型的預(yù)訓(xùn)練，無需依賴模型并行、檢查點或卸載策略。

• 論文地址：https://arxiv.org/pdf/2412.05270
• 論文網(wǎng)站：https://zhuhanqing.github.io/APOLLO/
• 論文代碼: https://github.com/zhuhanqing/APOLLO

方法簡要介紹

從 Adam 到結(jié)構(gòu)化更新

在傳統(tǒng)的 Adam 優(yōu)化器中，學(xué)習(xí)率的自適應(yīng)更新是基于每個參數(shù)的梯度動量和方差：

APOLLO 首先將 Adam 的更新過程等效理解為對原始梯度  施加一個基于梯度動量和方差的梯度縮放因子，進一步發(fā)現(xiàn)這一縮放因子可以通過結(jié)構(gòu)化更新（例如按通道或張量進行縮放）來近似實現(xiàn)。

在 LLaMA-130M 上的實驗驗證表明，結(jié)構(gòu)化更新不僅不會導(dǎo)致性能損失，反而會略微提升模型性能（見圖 1）。這一發(fā)現(xiàn)為 APOLLO 的設(shè)計提供了理論基礎(chǔ)。

圖 1: 結(jié)構(gòu)化梯度縮放與 Adam 的訓(xùn)練損失對比

APOLLO 框架

盡管結(jié)構(gòu)化更新簡化了優(yōu)化器設(shè)計，但仍需保存完整的優(yōu)化器狀態(tài)，內(nèi)存開銷依然較高。這促使解決以下關(guān)鍵問題：

問題①：是否可以將結(jié)構(gòu)化的學(xué)習(xí)率自適應(yīng)轉(zhuǎn)換為一種既實用又內(nèi)存高效的優(yōu)化方法？

作者提出了 APOLLO，在低秩空間中構(gòu)造一個輔助優(yōu)化器狀態(tài)，用于預(yù)測結(jié)構(gòu)化更新的梯度縮放因子（見圖 2）。由于優(yōu)化器狀態(tài)建立在低秩空間中，其內(nèi)存占用顯著減少。

圖 2: APOLLO 框架

問題②：低秩空間中的自適應(yīng)學(xué)習(xí)率能否有效近似其在原始空間中的行為？

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

這使得 APOLLO 可通過放縮縮放因子來彌補低秩帶來的誤差， 從而可采用極低的秩近似，在不犧牲性能的前提下實現(xiàn)極低的內(nèi)存消耗。實驗表明，在 LLaMA-7B 上，APOLLO 僅需 256 的秩，性能仍優(yōu)于使用 1024 秩的 GaLore。

圖 3：梯度縮放因子隨秩的變化按照理論推導(dǎo)的 √(n/r) 變化

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

不需要。APOLLO 無需傳統(tǒng)方法中基于昂貴 SVD 運算的低秩投影方法，只需采用基于隨機投影的方法。這種方法避免了 SVD 的計算開銷，引入了極低的額外計算，甚至 APOLLO 在 7B 上的優(yōu)化速度還快于 Adam。

問題④：是否能進一步壓縮優(yōu)化器狀態(tài)至 SGD 的內(nèi)存消耗，同時達到甚至超越 AdamW 的性能？

為實現(xiàn)極致內(nèi)存高效，APOLLO 提出了 APOLLO-Mini。APOLLO-Mini 通過僅使用 rank-1 的輔助子空間進行張量級梯度縮放，而非 APOLLO 中的通道級更新。

實驗結(jié)果顯示，APOLLO-Mini 在實現(xiàn) SGD 內(nèi)存消耗的同時，其預(yù)訓(xùn)練性能甚至超越了 AdamW。這一結(jié)果充分證明了結(jié)構(gòu)化學(xué)習(xí)率更新的有效性，并為資源有限的研究者提供了高效訓(xùn)練大模型的解決方案。

下圖為將 APOLLO 應(yīng)用于 Adam 的算法：

實驗及結(jié)果

研究者對 APOLLO 的 LLM 預(yù)訓(xùn)練、微調(diào)以及系統(tǒng)級性能優(yōu)勢進行了全面評估。

預(yù)訓(xùn)練

研究團隊在 C4 數(shù)據(jù)集（經(jīng)過凈化處理的 Common Crawl 語料庫）上對基于 LLaMA 的大型語言模型進行了預(yù)訓(xùn)練實驗。

實驗結(jié)果顯示 APOLLO 系列方法（即使在使用極低秩的情況下，如秩分別為 256 和 1）均能顯著降低內(nèi)存需求，并且獲得比 Adam 以及當(dāng)前 SOTA 方法更好的結(jié)果。

表 1: 使用 C4 數(shù)據(jù)集預(yù)訓(xùn)練不同大小的 LLaMA 模型驗證集困惑度（↓）對比

在 LLaMA 7B 模型上，APOLLO 系列不僅在內(nèi)存使用上取得巨大突破，其訓(xùn)練性能也遠超過 Adam 以及使用高秩（例如 1024）的 GaLore 方法。

表 2: 使用 C4 數(shù)據(jù)集預(yù)訓(xùn)練 LLaMA-7B 模型驗證集困惑度（↓）對比

微調(diào)

微調(diào)任務(wù)（常識推理和 MMLU）中，APOLLO 的表現(xiàn)同樣搶眼，達到了甚至超越 Adam 基線的效果。

系統(tǒng)優(yōu)勢

• 加速訓(xùn)練

在硬件資源受限的情況下，APOLLO 能夠支持比 AdamW 高達 4 倍的批量大小，從而實現(xiàn)最高可達 3 倍的訓(xùn)練加速。

• 極低內(nèi)存消耗

首次實現(xiàn)了在 12GB 內(nèi)存的消費級 GPU 上成功預(yù)訓(xùn)練 LLaMA 7B 模型，為大模型訓(xùn)練在低資源環(huán)境下提供了新的可能性。

圖 4：APOLLO 內(nèi)存和訓(xùn)練速度優(yōu)勢

• 極低的計算開銷

由于不依賴昂貴的 SVD 計算，APOLLO 優(yōu)化時間在 7B 上甚至比 adam 還要快。

隨著大語言模型規(guī)模不斷攀升，內(nèi)存需求也在呈爆炸式增長。APOLLO 不僅在理論上打破了優(yōu)化器內(nèi)存瓶頸，更在實踐中展現(xiàn)出預(yù)訓(xùn)練與微調(diào)的卓越性能。APOLLO 為資源受限的團隊提供了一條低成本、高效率的大模型訓(xùn)練之路，預(yù)示著未來大模型訓(xùn)練將進入一個全新的時代。