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

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

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

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

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

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

接下來我們看看文章主要內容。

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

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

為了解決上述問題，研究人員不斷開發各種優化技術，以減少預訓練和微調期間的內存使用。

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

看到這項研究后，網友表示：「是時候忘記云、忘記 HPC 了，有了 GaLore，所有的 AI4Science 都將在 2000 美元的消費級 GPU 上完成。」

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

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

因此，與 LoRA 不同，GaLore 不會改變訓練動態，可用于從頭開始預訓練 7B 模型，無需任何內存消耗的預熱。GaLore 也可用于進行微調，產生與 LoRA 相當的結果」。

方法介紹

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

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

本文證明了 GaLore 在預訓練和微調方面表現良好。當在 C4 數據集上預訓練 LLaMA 7B 時，8-bit GaLore 結合了 8-bit 優化器和逐層權重更新技術，實現了與全秩相當的性能，并且 optimizer state 的內存成本不到 10%。

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

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

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

實驗及結果

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

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

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

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

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

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

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

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

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

更多技術細節，請閱讀論文原文。