上個(gè)月，Meta FAIR 田淵棟參與的一項(xiàng)研究廣受好評(píng)，他們?cè)谡撐摹?nbsp;MobileLLM: Optimizing Sub-billion Parameter Language Models for On-Device Use Cases》中開(kāi)始卷 10 億以下參數(shù)小模型，主打在移動(dòng)設(shè)備上運(yùn)行 LLM。

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

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

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

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

論文標(biāo)題：GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection

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

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

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

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

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

看到這項(xiàng)研究后，網(wǎng)友表示：「是時(shí)候忘記云、忘記 HPC 了，有了 GaLore，所有的 AI4Science 都將在 2000 美元的消費(fèi)級(jí) GPU 上完成。」

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

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

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

方法介紹

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

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

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

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

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

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

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

研究者對(duì) GaLore 的預(yù)訓(xùn)練和 LLM 的微調(diào)進(jìn)行了評(píng)估。所有實(shí)驗(yàn)均在英偉達(dá) A100 GPU 上進(jìn)行。

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

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

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

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

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

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

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

如圖 4 所示，與 BF16 基準(zhǔn)和 8  bit Adam 相比，8 bit GaLore 所需的內(nèi)存要少得多，在預(yù)訓(xùn)練 LLaMA 7B 時(shí)僅需 22.0G 內(nèi)存，且每個(gè) GPU 的 token 批量較小（最多 500 個(gè) token）。

更多技術(shù)細(xì)節(jié)，請(qǐng)閱讀論文原文。