大模型高效微調-LoRA原理詳解和訓練過程深入分析

部落格首發於我的知乎，詳見：https://zhuanlan.zhihu.com/p/702629428

一、LoRA原理

LoRA(Low-Rank Adaptation of LLMs)，即LLMs的低秩適應，是引數高效微調最常用的方法。

LoRA的本質就是用更少的訓練引數來近似LLM全引數微調所得的增量引數，從而達到使用更少視訊記憶體佔用的高效微調。

1.1 問題定義

LoRA與訓練目標是解耦的，但本文設定就是語言模型建模。

以下將給出語言建模（可自然推廣到序列建模）的基本符號定義，即最大化給定提示的條件機率（本質是極大似然估計）。

The maximization of conditional probabilities given a task-specific prompt

給定一個引數為\(\mathbf{\Phi}\)預訓練的自迴歸語言模型$ P_{\Phi}(y|x)$。

\(x\)為輸入，\(y\)為輸出

note: 為與原文符號一致，下文\(\mathbf{\Phi}\)、\(\mathbf{\Theta}\)、\(\mathbf{W}\)均表示模型引數

全引數微調

每次full fine-tuning訓練，學一個 \(\Delta \mathbf{\Phi}\)，\(|\Delta \mathbf{\Phi}|\) 引數量大hold不住

語言模型的條件機率分佈建模目標

高效微調

$ \Delta \mathbf{\Phi}$ 是特定於下游任務的增量引數

LoRA將 $ \Delta \mathbf{\Phi}=\Delta \mathbf{\Phi}(\Theta)$ ，用引數量更少的$ \mathbf{\Theta}$來編碼（低秩降維表示來近似）, \(|\mathbf{\Phi}| << | \mathbf{\Theta}|\)

LoRA訓練目標

Transformer架構引數

Transformer層的輸入和輸出維度大小 \(d_{model}\)

\(\mathbf{W_q}\)、\(\mathbf{W_k}\)、\(\mathbf{W_v}\),和\(\mathbf{W_o}\)分別代表自注意力的query、key、value和output投影矩陣

\(\mathbf{W}\)或\(\mathbf{W}_0\)代表預訓練的權重矩陣

\(∆\mathbf{W}\)是微調後得到的增量引數矩陣（訓練後，最佳化演算法在引數上的累計更新量）

\(r\)代表LoRA模組的秩

1.2 LoRA簡介

LoRA的核心思想是，在凍結預訓練模型權重後，將可訓練的低秩分解矩陣注入到的Transformer架構的每一層中，從而大大減少了在下游任務上的可訓練引數量。

LoRA結構

We propose Low-Rank Adaptation(LoRA), which freezes the pre trained model weights and injects trainable rank decomposition matrices into each layer of the Transformer architecture, greatly reducing the number of trainable parameters for downstream tasks.

在推理時，對於使用LoRA的模型來說，可直接將原預訓練模型權重與訓練好的LoRA權重合並，因此在推理時不存在額外開銷。

1.3 為什麼要LoRA

背景

通常，凍結預訓練模型權重，再額外插入可訓練的權重是常規做法，例如Adapter。可訓練的權重學習的就是微調資料的知識。

但它們的問題在於，不僅額外增加了引數，而且還改變了模型結構。

這會導致模型訓練、推理的計算成本和記憶體佔用急劇增加，尤其在模型引數需在多GPU上分散式推理時（這越來越常見）。

推理效能比較

動機

深度網路由大量Dense層構成，這些引數矩陣通常是滿秩的。

相關工作表明，When adapting to a specific task, 訓練學到的過度引數化的模型實際上存在於一個較低的內在維度上（高維資料實際是在低維子空間中）

We take inspiration from Li et al. (2018a); Aghajanyan et al. (2020) which show that the learned over-parametrized models in fact reside on a low intrinsic dimension.

低秩矩陣

LoRA就假設LLM在下游任務上微調得到的增量引數矩陣\(\Delta \mathbf{W}\)是低秩的（肯定不是滿秩），即存在冗餘引數或高度相關的引數矩陣，但實際有效引數是更低維度的。

We hypothesize that the change in weights during model adaptation also has a low “intrinsic rank”, leading to our proposed Low-Rank Adaptation (LoRA) approach.

LoRA遂設想，對全引數微調的增量引數矩陣\(\Delta \mathbf{W}\)進行低秩分解近似表示（即對引數做降維）。

PCA降維示意圖，源於https://lightning.ai/pages/community/tutorial/lora-llm/

這樣訓練\(\Delta \mathbf{W}\)的低秩分解近似引數矩陣，效果上相比其他PEFT方法不會打什麼折扣，而且還能在推理時不增加額外開銷。

LoRA allows us to train some dense layers in a neural network indirectly by optimizing rank decomposition matrices of the dense layers’ change during adaptation instead, while keeping the pre-trained weights frozen

LoRA的大體思路就是這樣，具體的矩陣分解也是靠微調過程學習的。

接下來，介紹LoRA的具體方案。

1.4 LoRA實現

LoRA就是低秩矩陣適應，在凍結原有LLM引數時，用引數量更小的矩陣進行低秩近似訓練。

LoRA原理

對於預訓練權重矩陣\(\mathbf{W}_{0} \in \mathbb{R}^{d \times d}\)，LoRa限制了其更新方式，即將全參微調的增量引數矩陣\(\Delta \mathbf{W}\)表示為兩個引數量更小的矩陣$\mathbf{B} \(和\)\mathbf{A}$的低秩近似：

\[\mathbf{W}_{0} + \Delta \mathbf{W} = \mathbf{W}_{0}+ \mathbf{B}\mathbf{A} \]

其中，\(\mathbf{B}\in \mathbb{R}^{d \times r}\)和\(\mathbf{A}\in \mathbb{R}^{r \times d}\)為LoRA低秩適應的權重矩陣，秩\(r\)遠小於\(d\)。

此時，微調的引數量從原來\(\Delta \mathbf{W}\)的\(d*d\)，變成了\(\mathbf{B}\) 和 \(\mathbf{A}\)的\(2*r*d\)。可知，\(2*r*d << d*d\)（有\(2r << d\)）

給定輸入\(\mathbf{x} \in \mathbb{R}^{d}\)，新增LoRA後的輸出\(\mathbf{h} \in \mathbb{R}^{d}\)：

\[\mathbf{h} = (\mathbf{W}_{0} + \Delta \mathbf{W} ) \mathbf{x} = \mathbf{W}_{0}\mathbf{x} + \mathbf{B}\mathbf{A} \mathbf{x} \]

這裡，將\(\Delta \mathbf{h}=\mathbf{B}\mathbf{A} \mathbf{x}\)，便於後續求導計算。

在訓練時，原始引數\(\mathbf{W}_{0}\)被凍結，意味著\(\mathbf{W}_{0}\)雖然會參與前向傳播和反向傳播，但是不會計算其對應梯度\(\frac{\partial L}{\partial \mathbf{W}_0}\)，更不會更新其引數。

在推理時，直接按上面的式子將\(\mathbf{B}\mathbf{A}\)合併到\(\mathbf{W}_{0}\)中，因此相比原始LLM不存在推理延時。

1.5 LoRA引數初始化

在開始訓練時：

• 矩陣 \(\mathbf{B}\) 透過高斯函式初始化，\(b_i \sim N(0, {\sigma_b}^2)\)

• 矩陣 \(\mathbf{A}\) 為全零初始化，\(a_i = 0\)

這使得訓練開始前，LoRA的旁路\(\mathbf{B}\mathbf{A}=0\)，那麼微調就能從預訓練權重\(\mathbf{W}_{0}\)開始。

這樣就能和全引數微調時一樣，能有相同的開始。

這個策略要求，至少\(\mathbf{B}\) 和 \(\mathbf{A}\)中有一個被初始化為全0項。

但如果，全被初始化為0，\(\mathbf{B}\) 和 \(\mathbf{A}\)就訓不動了。因為，\(\mathbf{B}\) 和 \(\mathbf{A}\)全0時，處於鞍點，兩個權重的梯度也全為0😅
（\(\mathbf{B}\) 的梯度 \(\frac{\partial L}{\partial \mathbf{B}}\)依賴\(\mathbf{A}\)，\(\mathbf{A}\) 的梯度\(\frac{\partial L}{\partial \mathbf{A}}\)依賴\(\mathbf{B}\)，如果僅一項為0訓練是可以啟動的）

1.6 LoRA權重係數\(\frac{\alpha}{r}\)

實際實現時，\(\Delta \mathbf{W} = \mathbf{B}\mathbf{A}\)會乘以係數\(\frac{\alpha}{r}\)與原始預訓練權重合並\(\mathbf{W}_{0}\)，\(\alpha\)是一個超參：

\[\mathbf{h} = (\mathbf{W}_{0} + \frac{\alpha}{r} \Delta \mathbf{W})\mathbf{x} \]

直觀來看，係數\(\frac{\alpha}{r}\)決定了在下游任務上微調得到的LoRA低秩適應的權重矩陣\(\mathbf{B}\mathbf{A}\)佔最終模型引數的比例。

給定一個或多個下游任務資料，進行LoRA微調：

• 係數\(\frac{\alpha}{r}\)越大，LoRA微調權重的影響就越大，在下游任務上越容易過擬合
• 係數\(\frac{\alpha}{r}\)越小，LoRA微調權重的影響就越小（微調的效果不明顯，原始模型引數受到的影響也較少）

一般來說，在給定任務上LoRA微調，讓\({\alpha}\)為\(r\)的2倍數。（太大學過頭了，太小學不動。）

根據經驗，LoRA訓練大概很難注入新的知識，更多是修改LLM的指令尊隨的能力，例如輸出風格和格式。原始的LLM能力，是在預訓練是獲得的（取決於引數量、資料規模X資料質量）。

LoRA的秩\(r\)決定，LoRA的低秩近似矩陣的擬合能力，實際任務需要調參挑選合適的秩\(r\)維度。係數\(\frac{\alpha}{r}\)中\(\alpha\)決定新老權重的佔比。

1.7 LoRA的秩\(r\)如何選擇

和推薦系統中的評分矩陣分解、文字的非負矩陣分解，以及奇異值分解一樣。LoRA的低秩分解近似矩陣\(\mathbf{B}\) 和 \(\mathbf{A}\)的秩\(r\)的大小，決定了其擬合能力。

理想的情況是找到一個秩\(r\)，使得LoRA的低秩近似結構\(\mathbf{B}\mathbf{A}\)能具備全引數微調的增量矩陣\(\Delta \mathbf{W}\) 的表達能力，能越接近越好。

秩\(r\)成為了LoRA的超引數，隨著秩\(r\)維度的不斷增加，參與訓練的引數量也隨之增加，LoRA的低秩適應能力將逐漸提高甚至過擬合。

論文基於GPT-3 175B，WikiSQL和MultiNLI資料上，進行了關於LoRA秩\(r\)選取的實驗分析

Weight Type指明對Attention的那部分引數做了低秩適應。可以發現，在這個2個資料集上，\(r\)=4,8時基本上要略優於\(r\)=64的效果。更高的\(r\)不一定帶來更好的效果。

作者指出，增加\(r\)並不能涵蓋更有意義的子空間，這表明低秩適應矩陣就足夠了。但是，並不可能期望一個小的\(r\)適用於每個任務或資料集

一些秩\(r\)選取經驗：

• 微調的下游任務：
  簡單任務所需的秩\(r\)不大，任務越難/多工混合的情況，需要更大的秩\(r\)

• 基座能力：
  越強的基座，所需的秩\(r\)應該更小。例如Qwen2-72B-Instruct對比Qwen2-7B-Instruct。

越強的基座在處理同等任務時，需要微調的樣本數也通常會更少些。

• 資料規模：
  資料規模越大，需要更大的秩\(r\)

1.8 LoRA的微調的引數選取

LoRA原始論文只研究了注意力引數\(\mathbf{W_q}\)、\(\mathbf{W_k}\)、\(\mathbf{W_v}\),和\(\mathbf{W_o}\)。

論文基於GPT-3 175B，對比分析了訓練預算有限時，關於LoRA的微調注意力引數的選擇

在訓練預算為18M時 (roughly 35MB if stored
in FP16) on GPT-3 175B，注意力權重全部選擇時的效果最佳。

這表明，即使全部的注意力引數即使秩更小時（\(r=2\)），相比秩更大的（\(r=8\)）部分注意力引數，具有更強的建模能力。

在實際中，一般會把FFN的引數也考慮進來。

二、LoRA訓練

LoRA反向傳播的過程

LoRA訓練時，將凍結預訓練權重 \(\mathbf{W_0}\)
，只最佳化低秩矩陣 \(\mathbf{B}\)和\(\mathbf{A}\)。

LoRA訓練後，只需儲存低秩矩陣的\(\mathbf{B}\)和\(\mathbf{A}\)引數。

2.1 LoRA訓練的梯度計算

LoRa的計算圖和梯度計算

\(\mathbf{B}\) 和 \(\mathbf{A}\)的梯度計算, \(\mathbf{W}_{0}\)不參與計算。

\[\frac{\partial L}{\partial \mathbf{B}} = \frac{\partial L}{\partial \mathbf{h}} \frac{\partial \mathbf{h}}{\partial \Delta \mathbf{h}} \frac{\partial \Delta \mathbf{h}}{\partial \mathbf{B}} = \frac{\partial L}{\partial \mathbf{h}} (\mathbf{Ax})^{\mathsf{T}} \in \mathbb{R}^{d \times r} \]

\[\frac{\partial L}{\partial \mathbf{A}} = \frac{\partial L}{\partial \mathbf{h}} \frac{\partial \mathbf{h}}{\partial \Delta \mathbf{h}} \frac{\partial \Delta \mathbf{h}}{\partial \mathbf{A}} = \mathbf{B}^{\mathsf{T}} \frac{\partial L}{\partial \mathbf{h}} \mathbf{x}^{\mathsf{T}} \in \mathbb{R}^{r \times d} \]

繼續回傳的梯度，包括\(\mathbf{W}_{0}\)這一路：

\[\frac{\partial L}{\partial \mathbf{x}} = \frac{\partial L}{\partial \mathbf{h}} \frac{\partial \mathbf{h}}{\partial \mathbf{x}} = \frac{\partial L}{\partial \mathbf{h}} (\mathbf{W}_{0} + \mathbf{B}\mathbf{A})^{\mathsf{T}} \in \mathbb{R}^{d} \]

2.2 反向傳播計算量

全量微調前向計算：$ \mathbf{h} = \mathbf{W}_0
\mathbf{x}$

全量微調反向計算：

\[\frac{\partial L}{\partial \mathbf{W}_{0}} = \frac{\partial L}{\partial \mathbf{h}} \frac{\partial \mathbf{h}}{\partial \mathbf{W}_{0}} = \frac{\partial L}{\partial \mathbf{h}} \mathbf{x}^{\mathsf{T}} \in \mathbb{R}^{d \times d} \]

\[\frac{\partial L}{\partial \mathbf{x}} = \frac{\partial L}{\partial \mathbf{h}} \frac{\partial \mathbf{h}}{\partial \mathbf{x}} = \frac{\partial L}{\partial \mathbf{h}} \mathbf{W}_{0}^{\mathsf{T}} \in \mathbb{R}^{d } \]

LoRA微調計算：$\mathbf{h} = \mathbf{W}_{0}\mathbf{x} + \mathbf{B}\mathbf{A} \mathbf{x} $

此時，微調的引數量從原來\(\Delta \mathbf{W}\)的\(d*d\)，變成了\(\mathbf{B}\) 和 \(\mathbf{A}\)的\(2*r*d\)。可知，\(2*r*d << d*d\)（有\(2r << d\)）

不考慮pytorch、llm.c或8-bit最佳化器、Float8的訓練最佳化實現。可以看到，光梯度計算的話，實際計算量是增加了的。

2.3 LoRA在哪裡減少了視訊記憶體佔用

\(\mathbf{B}\)和\(\mathbf{A}\)的梯度是小頭這裡，暫時忽略。

預訓練權重 \(\mathbf{W_0}\)的梯度儲存開銷，實際就是LoRA能大大減少了視訊記憶體佔用的關鍵。

在LoRA訓練時，\(\mathbf{W_0}\)仍然會參與前向傳播和反向傳播，但是不會計算其對應梯度\(\frac{\partial L}{\partial \mathbf{W}_0}\)，更不會更新其引數。

因此，這一步不再需要計算和儲存梯度\(\frac{\partial L}{\partial \mathbf{W}_0}\)，以及更新\(\mathbf{W_0}\)。

以\(d=4096, r=16\)為例，這部分減少的梯度視訊記憶體佔用粗略估計為：\(d*d - 2*d*r = 1 - \frac{2r}{d}\), 減少了99.2187%。

若以Adaw optimizer的視角來看，其最佳化器所需維護的states（梯度的一階矩（均值）和二階原始矩（有偏方差）），那麼視訊記憶體佔用減少地更多。

三、效率分析

按照LoRA論文報告的結果，LoRA微調使得在訓練GPT3 175B時的，視訊記憶體消耗從1.2TB降至350GB；

當\(r=4\)時，最終儲存的模型從350GB降至35MB，極大降低了訓練的開銷。

下表，來源於LlamaFactory Github展示的微調LLM的最小硬體依賴估算值。

方法	精度	7B	13B	30B	70B	8x7B	8x22B
Full	AMP	120GB	240GB	600GB	1200GB	900GB	2400GB
Full	16	60GB	120GB	300GB	600GB	400GB	1200GB
Freeze	16	20GB	40GB	80GB	200GB	160GB	400GB
LoRA/GaLore/BAdam	16	16GB	32GB	64GB	160GB	120GB	320GB
QLoRA	8	10GB	20GB	40GB	80GB	60GB	160GB
QLoRA	4	6GB	12GB	24GB	48GB	30GB	96GB
QLoRA	2	4GB	8GB	16GB	24GB	18GB	48GB

實際的使用情況：

• 一張16GB視訊記憶體 T4，僅夠6B或7B的模型在batchsize為1時，進行int4 QLoRA，這還只是考慮輸入輸出有限時。
• 一張32GB視訊記憶體 V100，大致夠6B或7B的模型在batchsize為1時，進行LoRA微調。
• 一張80GB視訊記憶體 A800，Qwen1.5 72B 進行int4 QLoRA，以及例如Baichuan13B / Qwen14B的LoRA微調
• 2張A800 80GB視訊記憶體，可以進行全參SFT或增量SFT

參考資料

猛猿：圖解大模型微調系列之：大模型低秩介面卡LoRA（原理篇）
LoRA 微調-MartinLwx's blog|
Parameter-Efficient LLM Finetuning With Low-Rank Adaptation (LoRA)

原創不易，轉載需註明出處！