只用 24G 視訊記憶體，消費級 GPU 就能搞定大模型了。
上個月，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）。
更多技術細節，請閱讀論文原文。