田徑棟等人新作：突破記憶體瓶頸，讓一塊4090預訓練7B大模型

Meta FAIR 田徑棟參與的研究計畫在上個月獲得了廣泛好評。在他們的論文《MobileLLM: Optimizing Sub-billion Parameter Language Models for On-Device Use Cases》中，他們開始探討如何優化10億以下參數的小型模型，旨在實現在行動裝置上運行大型語言模型的目標。

3 月 6 日，田徑棟團隊發布了最新的研究成果，這次專注於提高LLM記憶體的效率。除了田徑棟本人，研究團隊還包括來自加州理工學院、德州大學奧斯汀分校以及CMU的研究人員。這項研究旨在進一步優化LLM記憶體的效能，為未來的技術發展提供支援和指導。

他們共同提出了一種名為GaLore（Gradient Low-Rank Projection）的訓練策略，這種策略允許全參數學習，相較於LoRA 等常見的低秩自適應方法，GaLore具有更高的記憶體效率。

這項研究首次表明，在配備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）。

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

以上是田徑棟等人新作：突破記憶體瓶頸，讓一塊4090預訓練7B大模型的詳細內容。更多資訊請關注PHP中文網其他相關文章！