上個月，Meta FAIR 田淵棟參與的一項研究廣受好評，他們在論文《 MobileLLM: Optimizing Sub-billion Parameter Language Models for On-Device Use Cases》中開始卷 10 億以下參數(shù)小模型，主打在移動設備上運行 LLM。

3 月 6 日，田淵棟又一項研究出爐，這次，他們主攻 LLM 內(nèi)存效率。除了田淵棟本人，還有來自加州理工學院、德克薩斯大學奧斯汀分校以及 CMU 的研究者。

他們合作提出了 GaLore（Gradient Low-Rank Projection），這是一種允許全參數(shù)學習的訓練策略，但比 LoRA 等常見的低秩自適應方法具有更高的內(nèi)存效率。

該研究首次證明了在具有 24GB 內(nèi)存的消費級 GPU（例如 NVIDIA RTX 4090）上預訓練 7B 模型的可行性，無需模型并行、檢查點或卸載策略。

論文地址：https://arxiv.org/abs/2403.03507

論文標題：GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection

接下來我們看看文章主要內(nèi)容。

現(xiàn)階段，大語言模型（LLM）在多個學科中表現(xiàn)出令人深刻的印象，然而大家又不得不面對這樣一個問題，預訓練和微調(diào) LLM 不僅需要大量計算，而且還需要大量內(nèi)存。

LLM 對內(nèi)存的需求不僅包括以億計算的參數(shù)，還包括梯度和 Optimizer States（例如 Adam 中的梯度動量和方差），這些參數(shù)可能大于存儲本身。舉例來說，使用單個批大小且從頭開始預訓練的 LLaMA 7B ，需要至少 58 GB 內(nèi)存（14 GB 用于可訓練參數(shù)，42 GB 用于 Adam Optimizer States 和權重梯度，2 GB 用于激活）。這使得訓練 LLM 在消費級 GPU（例如具有 24GB 內(nèi)存的 NVIDIA RTX 4090）上變得不可行。

為了解決上述問題，研究人員不斷開發(fā)各種優(yōu)化技術，以減少預訓練和微調(diào)期間的內(nèi)存使用。

該方法在 Optimizer States 下將內(nèi)存使用量減少了 65.5%，同時還能保持在 LLaMA 1B 和 7B 架構上使用最多 19.7B token 的 C4 數(shù)據(jù)集進行預訓練的效率和性能，以及在 GLUE 任務上微調(diào) RoBERTa 的效率和性能。與 BF16 基準相比，8-bit GaLore 進一步減少了優(yōu)化器內(nèi)存達 82.5%，總訓練內(nèi)存減少了 63.3%。

看到這項研究后，網(wǎng)友表示：「是時候忘記云、忘記 HPC 了，有了 GaLore，所有的 AI4Science 都將在 2000 美元的消費級 GPU 上完成?！?/span>

田淵棟表示：「有了 GaLore，現(xiàn)在可以在具有 24G 內(nèi)存的 NVidia RTX 4090s 中預訓練 7B 模型了。

我們沒有像 LoRA 那樣假設低秩權重結構，而是證明權重梯度自然是低秩的，因此可以投影到（變化的）低維空間中。因此，我們同時節(jié)省了梯度、Adam 動量和方差的內(nèi)存。

因此，與 LoRA 不同，GaLore 不會改變訓練動態(tài)，可用于從頭開始預訓練 7B 模型，無需任何內(nèi)存消耗的預熱。GaLore 也可用于進行微調(diào)，產(chǎn)生與 LoRA 相當?shù)慕Y果」。

方法介紹

前面已經(jīng)提到，GaLore 是一種允許全參數(shù)學習的訓練策略，但比常見的低秩自適應方法（例如 LoRA）更節(jié)省內(nèi)存。GaLore 關鍵思想是利用權重矩陣 W 的梯度緩慢變化的低秩結構，而不是試圖將權重矩陣直接近似為低秩形式。

本文首先從理論上證明了梯度矩陣 G 在訓練過程中會變成低秩，在理論的基礎上，本文用 GaLore 來計算兩個投影矩陣 和將梯度矩陣 G 投影成低秩形式 P^?GQ。在這種情況下，依賴于組件梯度統(tǒng)計的 Optimizer States 的內(nèi)存成本可以大大減少。如表 1 所示，GaLore 的內(nèi)存效率比 LoRA 更高。實際上，與 LoRA 相比，在預訓練期間，這可減少高達 30% 的內(nèi)存。

本文證明了 GaLore 在預訓練和微調(diào)方面表現(xiàn)良好。當在 C4 數(shù)據(jù)集上預訓練 LLaMA 7B 時，8-bit GaLore 結合了 8-bit 優(yōu)化器和逐層權重更新技術，實現(xiàn)了與全秩相當?shù)男阅?，并?optimizer state 的內(nèi)存成本不到 10%。

值得注意的是，對于預訓練，GaLore 在整個訓練過程中保持低內(nèi)存，而不需要像 ReLoRA 那樣進行全秩訓練。得益于 GaLore 的內(nèi)存效率，這是首次可以在具有 24GB 內(nèi)存的單個 GPU（例如，在 NVIDIA RTX 4090 上）上從頭開始訓練 LLaMA 7B，而無需任何昂貴的內(nèi)存卸載技術（圖 1）。

作為一種梯度投影方法，GaLore 與優(yōu)化器的選擇無關，只需兩行代碼即可輕松插入現(xiàn)有優(yōu)化器，如算法 1 所示。

下圖為將 GaLore 應用于 Adam 的算法：

實驗及結果

研究者對 GaLore 的預訓練和 LLM 的微調(diào)進行了評估。所有實驗均在英偉達 A100 GPU 上進行。

為了評估其性能，研究者應用 GaLore 在 C4 數(shù)據(jù)集上訓練基于 LLaMA 的大型語言模型。C4 數(shù)據(jù)集是 Common Crawl 網(wǎng)絡抓取語料庫的一個巨大的凈化版本，主要用于預訓練語言模型和單詞表征。為了最好地模擬實際的預訓練場景，研究者在不重復數(shù)據(jù)的情況下，對足夠大的數(shù)據(jù)量進行訓練，模型大小范圍可達 70 億個參數(shù)。

本文沿用了 Lialin 等人的實驗設置，采用了基于 LLaMA3 的架構，帶有 RMSNorm 和 SwiGLU 激活。對于每種模型大小，除了學習率之外，他們使用了相同的超參數(shù)集，并以 BF16 格式運行所有實驗，以減少內(nèi)存使用，同時在計算預算相同的情況下調(diào)整每種方法的學習率，并報告最佳性能。

此外，研究者使用 GLUE 任務作為 GaLore 與 LoRA 進行內(nèi)存高效微調(diào)的基準。GLUE 是評估 NLP 模型在各種任務中性能的基準，包括情感分析、問題解答和文本關聯(lián)。

本文首先使用 Adam 優(yōu)化器將 GaLore 與現(xiàn)有的低秩方法進行了比較，結果如表 2 所示。

研究者證明，GaLore 可以應用于各種學習算法，尤其是內(nèi)存高效的優(yōu)化器，以進一步減少內(nèi)存占用。研究者將 GaLore 應用于 AdamW、8 bit Adam 和 Adafactor 優(yōu)化器。他們采用一階統(tǒng)計的 Adafactor，以避免性能下降。

實驗在具有 10K 訓練步數(shù)的 LLaMA 1B 架構上對它們進行了評估，調(diào)整了每種設置的學習率，并報告了最佳性能。如圖 3 所示，下圖表明，GaLore 可適用于流行的優(yōu)化器，例如 AdamW、8-bit Adam 和 Adafactor。此外，引入極少數(shù)超參數(shù)不會影響 GaLore 的性能。

如表 4 所示，在大多數(shù)任務中，GaLore 都能以更少的內(nèi)存占用獲得比 LoRA 更高的性能。這表明，GaLore 可以作為一種全棧內(nèi)存高效訓練策略，用于 LLM 預訓練和微調(diào)。

如圖 4 所示，與 BF16 基準和 8  bit Adam 相比，8 bit GaLore 所需的內(nèi)存要少得多，在預訓練 LLaMA 7B 時僅需 22.0G 內(nèi)存，且每個 GPU 的 token 批量較?。ㄗ疃?500 個 token）。

更多技術細節(jié)，請閱讀論文原文。