全參數(shù)微調(diào)的顯存使用量和推理一樣多，大模型不再只是大型科技公司的玩具了。

在大模型方向上，科技巨頭在訓(xùn)更大的模型，學(xué)界則在想辦法搞優(yōu)化。最近，優(yōu)化算力的方法又上升到了新的高度。

大型語言模型（LLM）徹底改變了自然語言處理（NLP）領(lǐng)域，展示了涌現(xiàn)、頓悟等非凡能力。然而，若想構(gòu)建出具備一定通用能力的模型，就需要數(shù)十億參數(shù)，這大幅提高了 NLP 研究的門檻。在 LLM 模型調(diào)優(yōu)過程中通常又需要昂貴的 GPU 資源，例如 8×80GB 的 GPU 設(shè)備，這使得小型實(shí)驗室和公司很難參與這一領(lǐng)域的研究。

最近，人們正在研究參數(shù)高效的微調(diào)技術(shù)（PEFT），例如 LoRA 和 Prefix-tuning，為利用有限資源對 LLM 進(jìn)行調(diào)優(yōu)提供了解決方案。然而，這些方法并沒有為全參數(shù)微調(diào)提供實(shí)用的解決方案，而全參數(shù)微調(diào)已被公認(rèn)為是比參數(shù)高效微調(diào)更強(qiáng)大的方法。

在上周復(fù)旦大學(xué)邱錫鵬團(tuán)隊提交的論文《Full Parameter Fine-tuning for Large Language Models with Limited Resources》中，研究人員提出了一種新的優(yōu)化器 LOw-Memory Optimization（LOMO）。

通過將 LOMO 與現(xiàn)有的內(nèi)存節(jié)省技術(shù)集成，與標(biāo)準(zhǔn)方法（DeepSpeed 解決方案）相比，新方法將內(nèi)存使用量減少到了之前的 10.8%。因此，新方法能夠在一臺具有 8×RTX 3090 的機(jī)器上對 65B 模型進(jìn)行全參數(shù)微調(diào)，每個 RTX 3090 具有 24GB 內(nèi)存。

論文鏈接：https://arxiv.org/abs/2306.09782

在該工作中，作者分析了 LLM 中內(nèi)存使用的四個方面：激活、優(yōu)化器狀態(tài)、梯度張量和參數(shù)，并對訓(xùn)練過程進(jìn)行了三方面的優(yōu)化：

從算法的角度重新思考了優(yōu)化器的功能，發(fā)現(xiàn) SGD 在微調(diào) LLM 完整參數(shù)方面是一種很好的替代品。這使得作者可以刪除優(yōu)化器狀態(tài)的整個部分，因為 SGD 不存儲任何中間狀態(tài)。

新提出的優(yōu)化器 LOMO 將梯度張量的內(nèi)存使用量減少到 O （1），相當(dāng)于最大梯度張量的內(nèi)存使用量。

為了使用 LOMO 穩(wěn)定混合精度訓(xùn)練，作者集成了梯度歸一化、損失縮放，并在訓(xùn)練期間將某些計算轉(zhuǎn)換為全精度。

新技術(shù)讓內(nèi)存的使用等于參數(shù)使用加上激活和最大梯度張量。全參數(shù)微調(diào)的內(nèi)存使用被推向了極致，其僅等同于推理的使用。這是因為 forward+backward 過程的內(nèi)存占用應(yīng)該不會比單獨(dú)的 forward 過程少。值得注意的是，在使用 LOMO 節(jié)省內(nèi)存時，新方法確保了微調(diào)過程不受影響，因為參數(shù)更新過程仍然等同于 SGD。

該研究評估了 LOMO 的內(nèi)存和吞吐量性能，表明借助 LOMO，研究者在 8 個 RTX 3090 GPU 上就可以訓(xùn)練 65B 參數(shù)的模型。此外，為了驗證 LOMO 在下游任務(wù)上的性能，他們應(yīng)用 LOMO 來調(diào)優(yōu) SuperGLUE 數(shù)據(jù)集集合上 LLM 的全部參數(shù)。結(jié)果表明了 LOMO 對具有數(shù)十億參數(shù)的 LLM 進(jìn)行優(yōu)化的有效性。

方法介紹

在方法部分，本文詳細(xì)介紹了 LOMO（LOW-MEMORY OPTIMIZATION）。一般而言，梯度張量表示一個參數(shù)張量的梯度，其大小與參數(shù)相同，這樣一來內(nèi)存開銷較大。而現(xiàn)有的深度學(xué)習(xí)框架如 PyTorch 會為所有參數(shù)存儲梯度張量。現(xiàn)階段，存儲梯度張量有兩方面原因：計算優(yōu)化器狀態(tài)以及歸一化梯度。

由于該研究采用 SGD 作為優(yōu)化器，因此沒有依賴于梯度的優(yōu)化器狀態(tài)，并且他們有一些梯度歸一化的替代方案。

他們提出了 LOMO，如算法 1 所示，LOMO 將梯度計算與參數(shù)更新融合在一個步驟中，從而避免了梯度張量的存儲。

下圖為 SGD 和 LOMO 在反向傳播和參數(shù)更新階段的比較。Pi 為模型參數(shù)，Gi 為 Pi 對應(yīng)的梯度。LOMO 將梯度計算和參數(shù)更新融合到一個步驟中，使梯度張量最小。

LOMO 對應(yīng)的算法偽代碼：

具體而言，該研究將 vanilla 梯度下降表示為

，這是一個兩步過程，首先是計算梯度，然后更新參數(shù)。融合版本為

。 該研究的關(guān)鍵思想是在計算梯度時立即更新參數(shù)，這樣就不會在內(nèi)存中存儲梯度張量。這一步可以通過在向反向傳播中注入 hook 函數(shù)來實(shí)現(xiàn)。PyTorch 提供了注入 hook 函數(shù)的相關(guān) API，但卻無法用當(dāng)前的 API 實(shí)現(xiàn)精確的即時更新。相反，該研究在內(nèi)存中最多存儲一個參數(shù)的梯度，并隨著反向傳播逐一更新每個參數(shù)。本文方法減少了梯度的內(nèi)存使用，從存儲所有參數(shù)的梯度到只存儲一個參數(shù)的梯度。

大部分 LOMO 內(nèi)存使用與參數(shù)高效微調(diào)方法的內(nèi)存使用一致，這表明 LOMO 與這些方法相結(jié)合只會導(dǎo)致梯度占用內(nèi)存的輕微增加。這樣就可以為 PEFT 方法調(diào)優(yōu)更多的參數(shù)。

實(shí)驗結(jié)果

在實(shí)驗部分，研究者從三個方面評估了他們提出的方法，即內(nèi)存使用情況、吞吐量和下游性能。如果不作進(jìn)一步解釋，所有的實(shí)驗都是用 7B 到 65B 的 LLaMA 模型進(jìn)行的。

內(nèi)存使用情況

研究者首先剖析了，在不同設(shè)置下，訓(xùn)練期間的模型狀態(tài)和激活的內(nèi)存使用情況。如表 1 所示，與 AdamW 優(yōu)化器相比，LOMO 優(yōu)化器的使用導(dǎo)致內(nèi)存占用大幅減少，從 102.20GB 減少到 14.58GB；與 SGD 相比，在訓(xùn)練 LLaMA-7B 模型時，內(nèi)存占用從 51.99GB 減少到 14.58GB。內(nèi)存用量的大幅減少主要?dú)w因于梯度和優(yōu)化器狀態(tài)的內(nèi)存需求減少。因此，在訓(xùn)練過程中，內(nèi)存大部分被參數(shù)占據(jù)，與推理過程中的內(nèi)存用量相當(dāng)。

如圖 2 所示，如果采用 AdamW 優(yōu)化器進(jìn)行 LLaMA-7B 訓(xùn)練，相當(dāng)大比例的內(nèi)存（73.7%）被分配給優(yōu)化器狀態(tài)。用 SGD 優(yōu)化器替換 AdamW 優(yōu)化器可以有效減少優(yōu)化器狀態(tài)占用內(nèi)存的百分比，從而減輕 GPU 內(nèi)存使用（從 102.20GB 減少到 51.99GB）。如果使用 LOMO，參數(shù)更新和 backward 會被融合到一個步驟中，進(jìn)一步消除優(yōu)化器狀態(tài)對內(nèi)存的需求。

吞吐量

研究者比較了 LOMO、AdamW 和 SGD 的吞吐性能。實(shí)驗是在一臺配備了 8 個 RTX 3090 GPU 的服務(wù)器上進(jìn)行的。

對于 7B 的模型，LOMO 的吞吐量呈現(xiàn)顯著優(yōu)勢，超過 AdamW 和 SGD 約 11 倍。這一重大改進(jìn)可歸功于 LOMO 在單個 GPU 上訓(xùn)練 7B 模型的能力，這減少了 GPU 間的通信開銷。與 AdamW 相比，SGD 的吞吐量略高，這可歸因于 SGD 排除了動量和方差的計算。

至于 13B 模型，由于內(nèi)存的限制，它無法在現(xiàn)有的 8 個 RTX 3090 GPU 上用 AdamW 訓(xùn)練。在這種情況下，模型的并行性對 LOMO 來說是必要的，LOMO 在吞吐量方面仍然優(yōu)于 SGD。這一優(yōu)勢歸功于 LOMO 的內(nèi)存高效特性，以及只需要兩個 GPU 以相同的設(shè)置來訓(xùn)練模型，從而降低了通信成本，提高了吞吐量。此外，在訓(xùn)練 30B 模型時，SGD 在 8 個 RTX 3090 GPU 上遇到了內(nèi)存不足（OOM）的問題，而 LOMO 在只有 4 個 GPU 的情況下表現(xiàn)良好。

最后，研究者使用 8 個 RTX 3090 GPU 成功訓(xùn)練了 65B 模型，實(shí)現(xiàn)了 4.93 TGS 的吞吐量。利用這樣的服務(wù)器配置和 LOMO，模型在 1000 個樣本上的訓(xùn)練過程（每個樣本包含 512 個 token）大約需要 3.6 小時。

下游性能

為了評估 LOMO 在微調(diào)大型語言模型方面的有效性，研究者進(jìn)行了一系列廣泛的實(shí)驗。他們將 LOMO 與其他兩種方法進(jìn)行比較，一種是不需要微調(diào)的 Zero-shot，另一種是目前很流行的參數(shù)高效微調(diào)技術(shù) LoRA。

表 3 結(jié)果顯示：

LOMO 的表現(xiàn)明顯好于 Zero-shot；

在大多數(shù)實(shí)驗中，LOMO 普遍優(yōu)于 LoRA；

LOMO 可以有效擴(kuò)展至 650 億參數(shù)的模型。

LOMO 和 LoRA 在本質(zhì)上是相互獨(dú)立的。為了驗證這一說法，研究者使用 LLaMA-13B 在 BoolQ 和 MultiRC 數(shù)據(jù)集上進(jìn)行了實(shí)驗。結(jié)果如圖 3 所示。

他們發(fā)現(xiàn)，LOMO 在持續(xù)增強(qiáng) LoRA 的性能，不管 LoRA 取得的結(jié)果有多高。這表明，LOMO 和 LoRA 采用的不同微調(diào)方法是互補(bǔ)的。具體來說，LOMO 專注于微調(diào)預(yù)訓(xùn)練模型的權(quán)重，而 LoRA 則調(diào)整其他模塊。因此，LOMO 不會影響到 LoRA 的性能；相反，它有助于對下游任務(wù)進(jìn)行更好的模型調(diào)優(yōu)。