為什麼要微調#
選擇 LLM 完成一個 NLP 任務,如何下手?
從下圖中就能很好的明白哪個操作適合完成你當前的任務:
如果你有時間和大量數據,你完全可以重新訓練模型;一定量的數據,可以對預訓練模型進行微調;數據不多,最好的選擇是 “in context learning”,上下文學習,如 RAG。
當然,這裡我們主要研究微調這部分,微調讓我們無需重新訓練模型就可以取得比原先更好的表現。
怎麼微調#
眾所周知,顯卡(顯存)是限制我們平民玩家玩 LLM 的瓶頸,大部分人只能購買消費級顯卡,如 RTX 系列,因此需要找到一種能夠利用這 16GB 顯存進行微調的聰明辦法。
微調的瓶頸#
在訓練 llama 7B 這樣中等參數量的模型時,我們可能需要大概 28GB VRAM 來存儲模型原先的參數(下面會講這是怎麼估計的),再需要等量的顯存來存儲訓練過程中的梯度,通常還需要用參數量兩倍的量來跟蹤優化器的狀態。
來計算一下:
那我缺的這塊 96GB 的顯存誰給我補啊?
解決問題#
半精度#
第一步就是加載模型本身,對於 7B 模型,其中的每個參數的單位都是 32 位浮點數。
一個字節是 8bits,所以 32 位需要 4 字節 (4B)。7Billion,也就是 70 億,總共需要的存儲大小就是。
()
我們這裡就超出了 28-16=12GB,因此需要想辦法將模型的參數打包成一個更小的形式,一個非常自然的想法就是對參數的單位下手,能不能改用 16 或者是 8 位的浮點數呢,(分別對應 2B、1B 的單位存儲空間),只要換成 F16 就可以把這部分的顯存需求折半。作為權衡,相對應的浮點數精度、表示範圍會降低,可能會出現梯度爆炸、梯度消失等情況。谷歌對此提出了 bfloat16 (brain float,腦浮點),其核心目的就是提供一種既能保持較寬的數值範圍 (相較於 IEEE 規範,exponent: 5bit 到 8bit),又能簡化硬件實現的浮點格式 (fraction: 10bit 到 7bit),從而在不犧牲過多精度的情況下加速深度學習模型的訓練和推理過程。
我們選擇 16 位浮點數,顯存需求減半後一張卡就夠用了:
量化#
這裡簡單說下神經網絡的訓練流程:
我們對輸入內容進行 forward pass (前向傳播),也就是激活,然後使用結果與預測目標進行比較,根據預測和實際目標之間的差異 (損失),計算損失函數對於每個參數的梯度 (偏導數) 用於 BP (反向傳播),選擇一個優化算法 (如 SGD,即隨機梯度下降) 來更新參數,多輪迭代後得到模型。
模型中的梯度通常與原始模型中的參數擁有同樣的數據類型,對於每個參數都有對應的梯度,因此在不考慮優化器的時候都需要兩倍參數量的顯存。
一般會採用 Quantized (量化) 的方法,我們可以選擇 8 位浮點數
圖源自Nvidia 博客
可以看到在量化過程中,數據的表示範圍會被壓縮,數據會壓縮而集中,每個參數之間的差異減小,這可能導致大量的信息損失。對於那些超出新表示範圍的異常值進行剪切,可以減少由於這些極端數值引起的量化誤差。
選擇 int8 量化後,我們把模型參數和梯度所需要的內存砍到了 14GB:
LoRA#
儘管我們做了這麼多努力,但優化器才是關鍵部分。
業界喜愛而常用的 Adam 優化器效果很好,但也有相當高的內存佔用,原因如下:
Adam 優化器在每一步迭代中會更新參數 $\theta$,使用的更新公式如下(不必深入理解數學):
- 計算梯度的一階矩估計(即梯度的均值)和二階矩估計(即梯度的未中心化的方差)
其中,是在時間步的梯度,和是衰減率,通常取值接近 1,對應Karpathy 的 Batch-Norm 教程在第 3 小節介紹的指數移動平均。
- 對和進行偏差校正,以修正它們的偏差向 0 的初始化偏差
- 使用修正後的一階矩估計和二階矩估計來更新參數:
其中,是學習率,是為了保持數值穩定而加入的一個很小的常數。
這個過程在每個時間步重複,直到模型的參數收斂或達到某個停止條件。
對於1.
中的 momentum vector (動量向量),和 variance (方差) 向量,它們都各自有 7B 的參數,這是前面提到的需要 2 倍參數量的原因。
這裡的解決方案就是 LoRA (Low-Rank Adaptation),低秩適應:
這項技術可以減少可訓練參數數量,達到了減少模型權重佔用空間,加快訓練速度的效果。在這個情景下,LoRA 顯著地減少了需要被優化器以及梯度追蹤的參數量,減少了訓練過程中需要的顯存。
LoRA 背後的關鍵思想是,當對 llama2 這樣的大模型進行微調時,你不需要對每個參數都微調(也就是全參數微調),因為通常有一些參數和層相較於其它的來說更重要,如負責注意力機制和確定序列中哪些 token 與其它 token 相關以及相關方式的部分會更重要。LoRA 會取出這些特定的參數,注入低秩矩陣,後面訓練和傳播、更新參數時,修改的只是這個輔助的低秩矩陣。
LoRA 中的 R 超參數,也就是秩是可以調整的。但一般在實踐中,可能 LoRA 選取的特定參數只占總體的 10% 以下。
對於 LoRA 參數,選擇精度較高的 fp16,而優化器狀態的單位則是 fp32,因此內存佔量在這裡是參數量的四倍。
但這裡還有一個問題,那就是激活部分。激活的前向傳播過程中的開銷,是神經網絡中最大層的大小 $\times$ 批次大小 batch size (一次更新多少樣本),這可能還會佔用 5G 的內存,仍然超出了我們的預算。
QLoRA#
那我們能不能使用 4bit 量化呢?
這就是 QLoRA 這篇論文提出的思想,其所做的就是通過paged atom(分頁原子)優化技術,在需要時讓優化器狀態的分頁內存移到 CPU 上,減少訓練期間訓練峰值的影響:
為此,引入了新的單位nf4
(normal float 4)。
又能省下一些顯存:
梯度累積#
最後的一個問題在於 Batch Size 的選取上。如果我們選擇一次更新很少的樣本,訓練過程中的方差會很大,極端情況是完全的 SGD (隨機梯度下降)。所以一般會選擇中間地帶,也就是步伐在大而平滑與小而急促之間的 sweet spot,這也是為什麼一般都選擇 23, 64, 128 作為 batch size。
但我們現在只能一次加載一個樣本,因此引出了 Gradient Accumulation 技術。
其關鍵思想在於,不增加額外內存開銷的情況下獲得使用更大 batch size 的訓練效果。
其中的操作:
-
分批處理:將較大的批量數據分成多個小批量(這些小批量的大小是基於可用內存資源確定的)。對於每個小批量:
- 執行前向傳播來計算損失。
- 執行反向傳播來計算當前小批量的梯度,但不立即更新模型參數。
-
梯度累積:將每個小批量計算得到的梯度累加到之前的梯度上,而不是立即用它們來更新參數。
-
參數更新:在處理完所有的小批量數據,並累積了足夠數量的梯度之後,使用累積的梯度來一次性更新模型的參數。
微調實戰: Mistral 7B#
QLoRA
16GB VRAM
Mixtral 8x7B (MoE)#
硬體需求:>=65GB VRAM
感謝閱讀,我會儘快更新微調實戰部分~