Last Updated on 2023-09-21 by Clay
前言
大模型的浪潮自從 2022 年 11 月 ChatGPT 的發布後便一發不可收拾，直到現在開源的大型語言模型（Large Language Model）的量級還在不斷增大，比方說 LLaMA-2-70B、以及 Falcon-180B 等等。
大型語言模型的性能自然是相當優秀的，可是往往需要耗費大量且價格昂貴的 GPU 記憶體，這使得一些邊緣運算裝置根本就不可能讓模型進行推理（inference） —— 更遑論訓練、微調自己的模型了。
而本篇我所分享的 QLoRA 技術，是一篇在 2023 年 5 月由華盛頓大學發表的論文，通過以下技術完成了降低記憶體成本的模型訓練及推理：
- 提出 4-bit NormalFloat (NF4) 的量化資料型態
- 提出『二次量化』（Double Quantization）的機制減少記憶體開銷
- 使用 NVIDIA GPU 中對於 GPU-CPU 之間的記憶體分頁傳輸功能，降低記憶體使用率峰值的報錯問題
- 使用 LoRA 技術來降低 LLM 的 fine-tuning 成本
總地來說，我個人認為這是一篇相當實用的技術性論文，官方所釋出的 GitHub 項目，也能很方便地遷移到開發人員自己的開發環境中。
截止至我寫下閱讀筆記的今天（2023/09/12），已經有了 7.5k 個 stars。
QLoRA 介紹
研究團隊提出了 QLoRA，一種能夠在單顆 48 GB 的 GPU 上 finetune 65B 量級模型的方式，同時保留了 16-bit (半精度) 的微調性能。
QLoRA 是透過 4-bit 量化的模型計算出梯度，並以此來進行低秩適配器（Low Rank Adapters, LoRA）的反向傳播更新權重。
這裡也簡單記錄一下 LoRA。
LoRA 是一種不用改變模型原始參數就微調語言模型的方法。
實際微調任務中，我們會事先在要訓練的特定模型層旁（比如原始 LoRA 論文中在注意力機制中只添加在 Wq 和 Wv 這兩層）添加一組低秩（low rank）的適配器（adapter），並且輸出維度要跟原先的模型層一模一樣。之後，把適配器設定成可訓練、原始模型參數則凍結不允許訓練。
這樣我們就可以在不影響推理速度、只增加一點點參數量的情況下，輕鬆完成大型語言模型的訓練。
也是因為參數量不多的關係，所以模型在下游任務中容易收斂；並且由於原始模型的參數權重不變，所以防止過擬合（overfitting）。
而 QLoRA 正是在 LoRA 的這種微調技術上加入原始模型的量化；研究團隊根據這種量化方式所微調出來的模型家族，則被取名為 Guanaco（原駝）。
Guanaco 在 Vicuna benchmark 上，擊敗了之前的許多開源模型，達到了 ChatGPT 的 99.3% 性能，並且這只是在一顆 GPU 上進行了 24 小時的 fine-tuning 而已。（不過這種評估方式，並非沒有人詬病，故先暫且觀望參考）
Vicuna benchmark，是一組具有挑戰性的多輪對話，由 GPT-4 自動進行評分的標準測試基準。在 2023 年 6 月被棄用並以 MT-Bench 取代。
QLoRA 所提出的此類節省記憶體需求的訓練/推理方式，在實驗中顯示是不犧牲性能的。下面是 QLoRA 最仰賴幾個重要技術：
4-bit NormalFlaot (NF4)
一種新的資料型態，是常態分佈權重（normal distributed weights）的理論最佳資訊表示法。
在模型並未實際做神經網路層的計算時，所有的參數都以 4-bit 儲存著，一旦該層要做計算時，則依照需求轉換回 BF16 或 FP16 進行計算。
而從 4-bit 轉換回 16-bit 時，對應回的是以下 16 （2^4）個數字:
[-1.0, -0.6961928009986877, -0.5250730514526367, -0.39491748809814453, -0.28444138169288635, -0.18477343022823334, -0.09105003625154495, 0.0, 0.07958029955625534, 0.16093020141124725, 0.24611230194568634, 0.33791524171829224, 0.44070982933044434, 0.5626170039176941, 0.7229568362236023, 1.0]
這 16 個數字的選擇直到現在仍困擾著我，翻過了 bitsandbytes 的 issue，發現有人詢問，但到我記錄筆記的這天仍然未見詳細的解釋。
如果有大神明白這個計算方式，還請不吝告知。感謝！
根據網路上找到的資源，顯然計算方式可以從 bitsandbytes 的原始碼中找到：
def create_normal_map(offset=0.9677083, use_extra_value=True):
if use_extra_value:
# one more positive value, this is an asymmetric type
v1 = norm.ppf(torch.linspace(offset, 0.5, 9)[:-1]).tolist()
v2 = [0]*(256-15) ## we have 15 non-zero values in this data type
v3 = (-norm.ppf(torch.linspace(offset, 0.5, 8)[:-1])).tolist()
else:
v1 = norm.ppf(torch.linspace(offset, 0.5, 8)[:-1]).tolist()
v2 = [0]*(256-14) ## we have 14 non-zero values in this data type
v3 = (-norm.ppf(torch.linspace(offset, 0.5, 8)[:-1])).tolist()
v = v1 + v2 + v3
values = torch.Tensor(v)
values = values.sort().values
values /= values.max()
assert values.numel() == 256return values
如果我們實際試跑這段程式，我們會看到確實被印出來的 16 個數字。(過程中需要 padding 的部分，被我直接取代成單一個 0，這樣才能完整表現出 16 個數值)
import torch
from scipy.stats import norm
def create_normal_map(offset=0.9677083, use_extra_value=true):
if use_extra_value:
# one more positive value, this is an asymmetric type
v1 = norm.ppf(torch.linspace(offset, 0.5, 9)[:-1]).tolist()
#v2 = [0]*(256-15) ## we have 15 non-zero values in this data type
v2 = [0]
v3 = (-norm.ppf(torch.linspace(offset, 0.5, 8)[:-1])).tolist()
else:
v1 = norm.ppf(torch.linspace(offset, 0.5, 8)[:-1]).tolist()
v2 = [0]*(256-14) ## we have 14 non-zero values in this data type
v3 = (-norm.ppf(torch.linspace(offset, 0.5, 8)[:-1])).tolist()
v = v1 + v2 + v3
values = torch.tensor(v)
values = values.sort().values
values /= values.max()
#assert values.numel() == 256return values
defmain():
for idx, value inenumerate(create_normal_map(), 1):
print(value)
if __name__ == "__main__":
main()
Output:
tensor(-1.) tensor(-0.6962) tensor(-0.5251) tensor(-0.3949) tensor(-0.2844) tensor(-0.1848) tensor(-0.0911) tensor(0.) tensor(0.0796) tensor(0.1609) tensor(0.2461) tensor(0.3379) tensor(0.4407) tensor(0.5626) tensor(0.7230) tensor(1.)
看起來是與論文中提供的數字相符的。那麼接下來困擾我的問題來了。
torch.linspace(offset, 0.5, 9) 是一段從 offset 到 0.5 之間取 9 個點的函式，這一步目前沒有問題；
norm.ppf() 是計算常態分佈的百分點函數（Percent-Point Function, PPF），這點也很容易理解。例如我們若輸入 norm.ppf(0.95)，會取得 1.64 左右的值。這意味著在常態分佈下，大約有 95% 的資料小於等於 1.64。
這裡必須說明的是，如果我們在 norm.ppf() 中輸入 0 或是 1，會得到 -inf 和 inf，畢竟常態分佈的尾部雖然趨近於零但是永不為零。
所以找一個 offset 作為邊界是合適且必須的。而根據網路上找到的資料來看，作者曾有做出解釋：
We want to find the quantiles which have equal area to the left and the right side of the quantile. This means, we do not start with the 0 or 1 quantile for the normal distribution, but with an offset quantile. This start position is called offset in the code snipped and is 1-1/(2*15). If we have an asymmetric data type, we have one side with spacing equivalent to 16 “halves” around each quantile and the other side with 15 halves. As such, the offset is on average (1-1/(2*15) + 1-1/(2*16))/2 = 0.9677083.
作者希望找到一個特定的 offset 來確保選定的分位數在左右兩側面積相等（這裡左右兩側我理解為正負，不知道有沒有想錯。或許是指每個分位數左右兩側？但這樣想感覺也不合理）。
而計算方式就是 (1-1/(2*15) + 1-1/(2*16))/2 = 0.9677083。
這個問題困擾我近兩個禮拜，我決定暫時放下這個困惑。畢竟從原始論文中的公式也代不出程式碼給的數值。或許接下來會嘗試在 GitHub 上發 issue。
另一種可能就是程式碼中的兩種情況（use_extra_value 是否為 True），分別是切割了 16 個數字跟 15 個數字，並且他們都共用了同一個 offset —— 所以 offset 是分別考慮了以上兩種情況，所以計算平均時的 15 和 16 的值是這麼來的？我並未確認這件事，並且暫時沒有任何方法可以證明我的種種假設。
有趣的是，從 4-bit 還原回的這 16 個數字（NF4 量化），在本篇研究工作中認為是資訊上的理論最佳實作，但在我找資料時，找到了另外一份研究工作 NF4 Isn’t Information Theoretically Optimal (and that’s Good)。
有興趣的話可以研究看看。
Double Quantization
這個技術的部分比 NF4 的還原值容易理解多了。
諮詢了做量化技術的朋友，說是常翻譯成二次量化，簡單來說就是量化模型參數時需要保存一個量化常數（quantization constant），而這個量化常數因為是 FP32，所以可以進一步量化這個量化常數。
一般的量化，可以參考以下公式（這裡舉的例子是從 FP32 轉為 INT8）：
之所以是 127，是因為扣掉符號位元（sign），2^7 表示為 128，但是需要保留給 ±inf，所以表示範圍落在 [-127, 127] 之間。
X 可以視為一整個量化的區塊（block），因為量化是一個區塊一個區塊做的，所以我們需要使用 round(127 * (X / absmax(X))) 來讓最大值落 在 127 這一邊界上。簡單來講就是縮放到 8-bit 所能表示的範圍內。
而 127 / absmax(X) 就被稱為量化常數（quantization constant），每一個量化區塊都必須儲存著自己的 32-bit 的量化常數 —— 因為反量化需要用到量化常數還原。
由於量化時有取 round 做近似，所以反量化（dequantization）時不可能無損還原，一定有精度的損失。
以上是一般量化的介紹。在 QLoRA 中提出的二次量化，指的就是針對『量化常數』的再次量化。
剛剛提到，一般儲存的量化常數是 32-bit，這無形中增加了記憶體開銷。以論文中舉的例子來看，假設一個量化區塊（block）是 64 個參數，那麼一個參數平均多花了 32 / 64 = 0.5 bits。
也就是說在我們一個參數使用 4-bit 儲存的情況下，額外的記憶體開銷是 0.5 / 4 = 12.5%，是相當可觀的。比方說如果模型量化後需要 10GB 的 VRAM，那麼就需要額外花 1.25GB 的 VRAM 來儲存量化常數。
所以研究團隊就針對這個量化常數，再多做一次量化。量化公式與上面相同。
假設我們把量化常數量化為 INT8，並且針對量化常數進行二次量化的量化區塊大小為 256，代表我們現在可以把一個參數的額外記憶體開銷寫成 8 / 64 + 32 / (64 * 256) = 0.127 bits。
也就是說額外的記憶體開銷為 0.127 / 4 = 3.175%。現在如果模型量化後需要 10GB VRAM，額外只需要花 0.3GB 來儲存『量化常數』和『量化常數的量化常數』即可。
二次量化的數學表達式如下。
不過直觀上，這樣做雖然節省了記憶體，但是可以合理想像模型的推理速度會變慢；因為量化的模型，是在對每一個模型層做推理時，才把原先用 4-bit 儲存的參數還原回 16-bit，以減少記憶體的開銷。
所以多做了一次轉換，顯然更花時間。
另外，從上方公式 (5)、(6) 可以明顯看出，要訓練的 LoRA 層是沒有被量化的。
Paged Optimizer
管理記憶體峰值時的操作。可以做到 GPU 跟 CPU 之間的分頁交換，可以在原先 GPU 會發生 OOM（out of memory）的時候先把一部分的待處理計算移動到 CPU，等到 GPU 計算完後，再傳遞回 GPU。
這部分我比較沒有研究；或許哪天研究 CUDA 函式庫或相關套件的的時候會再認真看。
統整摘要
看過了 QLoRA 所使用的相關技術，現在能很好地想像 QLoRA 的工作流程。
LoRA 跟全微調的差異在於訓練的不是整個模型、而是 LoRA 的適配器，並且這個適配器的神經網路乘開後，是能夠直接加回原始模型層中的。這樣訓練的好處上方已經簡單帶過，這裡便不再贅述。
而 QLoRA 則是把原始模型做了量化，使用了研究團隊提出的 NF4 資料型態來確保更好的 16-bit 還原值。
並且為了更好地節省記憶體，研究團隊還提出了二次量化的技術，進一步減少記憶體開銷。
模型層基本上會一直保持著 4-bit 的狀態，但是若遇到輸入需要進行推理時，會再反量化回 16-bit —— 若是 track 一下 bitsandbytes 的實現，會發現它只是在 layer 進行 forward 時，使用『量化常數』跟『量化常數』把 layer 的參數表示回 16-bit 並對輸入做推理，實際上並不是真的有還原回一整個 16-bit 表示的 layer。
而訓練過程中，還使用了 NVIDIA 支援的 Paged Optimizer 來防止 OOM。
正因為做了這麼多巧妙的設計，才得出論文一開始提出的結論 —— 可以在 48GB 的 GPU 上訓練 65B 的模型。
要知道，65B 的模型本來光是載入記憶體，就需要至少 250GB+ 的記憶體了。
論文實驗結果
當然，若是只有量化技術使得記憶體開銷減少、實際上模型很難微調的話仍然不是我們想要的；因為那意味著我們只能拿這種量化技術來做 inference。
所以論文中，當然提出了一些實驗的結果，來證明 QLoRA 訓練出來的模型跟全精度、半精度、混合精度或其他量化訓練方式的效果是相仿的。
使用感想
最近我也自己上手跑了 QLoRA 官方 repo 的 qlora.py 腳本，也認真地研究了下不到千行的程式碼。
架構還算單純，畢竟複雜的量化、反量化操作都被包進跟 transformers 整合的 bitsandbytes 了。
我嘗試訓練了直接使用 alpaca 資料集的模型（也就是全是默認值）、以及使用我跟同事工作時所整理的一份資料集來微調 llama-2-7b —— 體驗的效果都是相當不錯的，確實能看到 fine-tune 所帶來的模型進步、也能明顯看到節省的大量記憶體。
總之，QLoRA 真的是個很棒的訓練流程，接下來我還會受益於這份訓練腳本無數次吧！