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

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

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

例如，預(yù)訓(xùn)練 LLaMA 7B 的一個(gè)批次就消耗高達(dá) 58GB 內(nèi)存：14GB 用于模型參數(shù)，42GB 用于優(yōu)化器狀態(tài)和梯度。這巨大的 “隱形” 內(nèi)存開(kāi)銷(xiāo)迫使研究者不得不選用顯存更大的高端 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)新方法首次實(shí)現(xiàn)了以 SGD 級(jí)別的內(nèi)存成本訓(xùn)練大模型，同時(shí)保持了 AdamW 的卓越優(yōu)化性能（已在 LLaMA 7B 的預(yù)訓(xùn)練上驗(yàn)證）。

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

核心亮點(diǎn)

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

卓越系統(tǒng)性能

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

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

方法簡(jiǎn)要介紹

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

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

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

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

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

APOLLO 框架

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

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

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

圖 2: APOLLO 框架

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

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

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

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

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

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

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

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

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

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

實(shí)驗(yàn)及結(jié)果

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

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

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

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

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

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

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

微調(diào)

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

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

• 加速訓(xùn)練

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

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

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

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

• 極低的計(jì)算開(kāi)銷(xiāo)

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

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