天天看點

在消費級GPU調試LLM的三種方法:梯度檢查點,LoRA和量化

作者:deephub

LLM的問題就是權重參數太大,無法在我們本地消費級GPU上進行調試,是以我們将介紹3種在訓練過程中減少記憶體消耗,節省大量時間的方法:梯度檢查點,LoRA和量化。

梯度檢查點

梯度檢查點是一種在神經網絡訓練過程中使動态計算隻存儲最小層數的技術。

為了了解這個過程,我們需要了解反向傳播是如何執行的,以及在整個過程中層是如何存儲在GPU記憶體中的。

1、前向和後向傳播的基本原理

前向傳播和後向傳播是深度神經網絡訓練的兩個階段。

在消費級GPU調試LLM的三種方法:梯度檢查點,LoRA和量化

在前向傳遞過程中,輸入被矢量化(将圖像轉換為像素,将文本轉換為嵌入),并且通過一系列線性乘法和激活函數(如sigmoid或ReLU等非線性函數)在整個神經網絡中處理每個元素。

神經網絡的輸出,被稱為頭部,被設計用來産生期望的輸出,例如分類或下一個單詞預測。然後将矢量化的預測結果與預期結果進行比較,并使用特定的損失函數(如交叉熵)計算損失。

基于損失值,以最小化損失為目标更新每層的權值和偏差。這個更新過程從神經網絡的末端開始并向起點傳播。

上面就是一個簡單的過程,下面才是我們主要關注的:計算是如何存儲在記憶體中的。

2、減少存儲數量

一種簡單的方法是隻保留反向傳播所需的基本層,并在它們的使用完成後從記憶體中釋放它們。

在消費級GPU調試LLM的三種方法:梯度檢查點,LoRA和量化

從上圖可以看出,同時存儲在記憶體中的層的最大數量并不是最優的。是以我們需要找到一種方法,在保持反向傳播工作的同時,在記憶體中存儲更少的元素。

3、減少計算時間

減少記憶體占用的一種方法是在神經網絡開頭的反向傳播過程中重新計算每一層。

在消費級GPU調試LLM的三種方法:梯度檢查點,LoRA和量化

但是在這種情況下,計算時間會明顯增加,使得訓練在大模型的情況下不可行。

4、優化計算和記憶體梯度檢查點

該技術通過儲存“檢查點”以計算反向傳播期間“丢失”的層。 該算法不是從頭開始計算層,如前面的示例所示,而是從最近的檢查點開始計算。

在消費級GPU調試LLM的三種方法:梯度檢查點,LoRA和量化
在消費級GPU調試LLM的三種方法:梯度檢查點,LoRA和量化

平衡記憶體存儲和計算時間的最佳政策是設定O(sqrt(n))個檢查點,層數為n。這樣,一次反向傳播計算的額外計算次數将對應于一次額外的前反向傳播。

這種技術可以在較小的gpu上訓練較大的模型,但代價是需要額外的計算時間(約20%)。

5、如何實作梯度檢查點

transformer庫已經提供了梯度檢查點技術。

from transformers import AutoModelForCausalLM, TraininArguments

model = AutoModelForCausalLM.from_pretrained(

model_id,

use_cache=False, # False if gradient_checkpointing=True

**default_args

)

model.gradient_checkpointing_enable()

LoRA

LoRA是微軟團隊開發的一種技術,用于加速大型語言模型的微調。他們在GPT-3 175B上實施了這種方法,并大大減少了訓練參數的數量。

在消費級GPU調試LLM的三種方法:梯度檢查點,LoRA和量化

他們的方法當機預訓練模型的所有參數,并将新的可訓練參數嵌入到transformer架構中的特定子產品中,如注意力子產品(查詢、鍵、值,但也适用于其他子產品)。

為了實作這些擴充卡,他們利用線性層,如下面的等式所示,其中x (dimension: d)和h (dim: k)作為乘法前後的層,Wo作為預訓練的權重,B和A作為新的權重矩陣。

在消費級GPU調試LLM的三種方法:梯度檢查點,LoRA和量化

矩陣B和A的維數分别為(d × r)和(r × k),且r << min(d, k)。

也就是說在不使訓練過程複雜化的情況下,将新的密集層添加到現有的層上。在微調過程中,權重矩陣BA初始化為0,并遵循α/r的線性尺度,α為常數。當使用Adam算法優化權重時,α與學習率大緻相同。

對不同的LoRA配置進行了測試,論文得出的結果是,将r=8(或更高)應用于各種子產品的性能最好。

在消費級GPU調試LLM的三種方法:梯度檢查點,LoRA和量化

一旦對LoRA模型進行了微調,就可以将權重合并在一起以獲得單個模型,或者隻單獨儲存擴充卡,并将預訓練模型與現有模型分開加載。

在消費級GPU調試LLM的三種方法:梯度檢查點,LoRA和量化

Hugging Face開發的PEFT庫,可以利用LoRA技術。

from peft import LoraConfig, TaskType
lora_config = LoraConfig(
r=16,
lora_alpha=16,
target_modules=["query_key_value"]
lora_dropout=0.1,
bias="none",
task_type=TaskType.CAUSAL_LM,
)           

還可以針對transformer架構中的所有密集層:

# From https://github.com/artidoro/qlora/blob/main/qlora.py
def find_all_linear_names(args, model):
cls = torch.nn.Linear
lora_module_names = set()
for name, module in model.named_modules():
if isinstance(module, cls):
names = name.split('.')
lora_module_names.add(names[0] if len(names) == 1 else names[-1])           

然後就是将“初始化”擴充卡添加到預訓練模型中。

from transformers import AutoModelForCausalLM
from peft import get_peft_model
model = AutoModelForCausalLM.from_pretrained(model_id)
lora_model = get_peft_model(model, peft_config)
lora_model.print_trainable_parameters()           

訓練完成後,可以單獨儲存擴充卡,也可以将它們合并到模型中。

# Save only adapaters
lora_model.save_pretrained(...)
# Save merged model
merged_model = lora_model.merge_and_unload() 
merged_model.save_pretrained(...)           

量化

談到LoRA,我就還需要說一下量化。這兩種技術在論文QLORA得到了高效的融合,并且已經通過bitsandbytes、peft和accelerayte整合到了Hugging Face 的transformer中。

1、什麼是量化?

量化是一種技術,可以降低元素的精度,但不會失去元素的整體意義。例如在圖檔的情況下,量化包括減少像素的數量,同時保持圖像的一個體面的分辨率。

在消費級GPU調試LLM的三種方法:梯度檢查點,LoRA和量化

上圖肉眼基本看不出差別,但是存儲空間卻少了很多。在解釋量化之前,需要了解計算機如何表示數字的

2、浮點數基本原理

在消費級GPU調試LLM的三種方法:梯度檢查點,LoRA和量化

計算機是二進制的,這意味着它們隻通過0和1交換資訊。為了表示數字,科學家設計了一種稱為浮點格式的特殊系統,它允許計算機了解大範圍的數值。最常見的表示形式是單精度浮點格式,由32位組成(1位= 0或1)。

除此以外還存在各種格式,例如半精度(16位)或雙精度(64位)。簡而言之,使用的比特數越多,可以容納的數字範圍就越廣。

像GPT-3.5或Bloom-175B這樣的模型非常大。在FP32格式中,這将表示:

175*10⁹. 4位元組= 700Gb,半精度為350Gb,基本不可能加載到GPU記憶體中,那麼我們如何縮小這些模型呢?

3、從FP32到Int8

Int8表示[- 127,127]之間的任何數字。

我們想将一個浮點數向量簡化為Int8格式:

v = [-1.2, 4.5, 5.4, -0.1]

我們需要做的是定義v的最大值(這裡是5.4),并将所有數字縮放到Int8[- 127,127]的範圍内。是以需要計算系數

α = 127 / max(v) = 127 / 5.4 ~ 23.5

然後把v中的所有數乘以α,然後四舍五入,得到:

α.v = [-28, 106, 127, -2]

如果想去量化這個向量,隻需要做相反的操作,就能夠得到初始向量!

v = [-1.2, 4.5, 5.4, -0.1]

可以看到量化和反量化不會丢失任何資訊。但實際上在四舍五入每個值時确實會失去精度。然而,在這個特定的例子中差異并不大,因為我們決定隻用一個小數來表示數字,另外就是對于大模型來說,參數互相很大,之間也有關系,是以四舍五入的精度丢失不會對模型的結果産生很大的影響(是不産生很大影響,不是沒影響),為了節省記憶體丢失一些小小的精度還是可以接受的。

那麼,如果有異常值存在會發生什麼?假設我們現在有這個向量:

v’ = [-1.2, 70, 5.4, -0.1]

目前的最高數字是70,這可以被視為一個異常值。如果我們重制完全相同的過程,我們在量化之後得到:

de-quantized v’ = [-1.1, 70, 5.5, 0.0]

精度的損失開始出現了,讓如果我們将同樣的損失應用于由70億個參數組成的LLM:缺乏精度将在整個神經網絡中積累,導緻有意義的資訊完全丢失,并導緻純噪聲。而且我們現在使用的是8位格式,如果是4位甚至3位,結果會更糟,對吧。

但是大佬們找到了一種将量化應用于LLM的方法!

4、LLM.int8()使大規模量化成為可能

論文LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale 介紹了一種繞過此異常值問題的方法。

在消費級GPU調試LLM的三種方法:梯度檢查點,LoRA和量化

量化參數的完整性會導緻性能下降,而在矩陣乘法過程中使用量化,結合混合精度分解和向量量化。在矩陣乘法過程中,從權重矩陣中提取包含異常值(高于門檻值)的向量,進而産生兩次乘法。 小數字矩陣(根據論文代表 99.9% 的值)被量化,而大數字則保留在 FP16 中。

按照混合精度分解原理,對小數乘法輸出進行反量化,并添加到其他輸出。

在消費級GPU調試LLM的三種方法:梯度檢查點,LoRA和量化

也就是說量化技術僅在推理(矩陣乘法)期間使用,這意味着實際上沒有8位數字組成的更小的模型!由于這種技術實作,我們甚至得到了一個更大的模型!(根據該論文,對于13B以下的模型,誤差為0.1%)但是在BLOOM-175B上的實驗表明,在沒有任何性能下降的情況下,記憶體占用減少了1.96倍!這種技術可以通路以前無法裝入GPU記憶體的大型模型

5、可以微調這個量化模型嗎?

不行,因為這種技術隻适用于推理,不适合訓練。

如果我們可以使用量化減少GPU記憶體占用,并使用LoRA技術訓練新的擴充卡,會怎麼樣?

還記得我們以前介紹的QLoRA嗎,它就幹的是這個事,他們成功地将預訓練模型量化為4位!它們通過一些新技術來成功地量化模型,比如雙量化和4位NormalFloat。

6、如何在代碼中使用量化?

首先需要安裝bitsandbytes和accelerate 庫

pip install -q bitsandbytes
pip install -q accelerate
pip install -q peft==0.4.1           

然後,在調用from_pretrained方法時,可以通過傳遞參數load_in_4bit=True或load_in_8bit= true來加載4位或8位量化的模型。

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("facebook/opt-350m", 
load_in_4bit=True, 
device_map="auto"
)           

也可以使用BitsAndBytesConfig類來進行進階的設定

from transformers import BitsAndBytesConfig
nf4_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_use_double_quant=True,
bnb_4bit_compute_dtype=torch.bfloat16
)
model_nf4 = AutoModelForCausalLM.from_pretrained(
model_id, 
device_map="auto"
quantization_config=nf4_config
)           

這樣模型差不多可以進行推斷了。但是我們還需要設定一下的參數:

當機量化參數以防止訓練,

在所有歸一化層和 LM 頭中使用FP32(未量化),以確定模型的穩定性

如果使用梯度檢查點,需要配置model.enable_input_require_grad()

for name, param in model.named_parameters():
# freeze base model's layers
param.requires_grad = False
# cast all non int8 or int4 parameters to fp32
for param in model.parameters():
if (param.dtype == torch.float16) or (param.dtype == torch.bfloat16):
param.data = param.data.to(torch.float32)
if use_gradient_checkpointing:
# For backward compatibility
model.enable_input_require_grads()           

在最新的peft==0.4.1庫中,使用prepare_model_for_kbit_training()方法可以處理這個準備工作。

這樣我們就有了一個量子的模型!

一段代碼總結

我們已經介紹了梯度檢查點、LoRA和量化,讓我們編寫代碼來對LLM進行微調。

先安裝必要的庫:

pip install -q -U bitsandbytes
pip install -q -U git+https://github.com/huggingface/transformers.git
pip install -q -U git+https://github.com/huggingface/peft.git
pip install -q -U git+https://github.com/huggingface/accelerate.git           

然後就是代碼:

from transformers import (
AutoModelForCausalLM,
BitsAndBytesConfig
)
from peft import (
get_peft_model,
LoraConfig,
TaskType,
prepare_model_for_kbit_training
)
# Import the model
gradient_checkpointing = True
model = AutoModelForCausalLM.from_pretrained(
args.model_id,
use_cache=False if gradient_checkpointing else True, # this is needed for gradient checkpointing
device_map="auto",
load_in_4bit=True
)
# Prepare the model (freeze, cast FP32, enable_require_grads, activate gradient checkpointing)
model = prepare_model_for_kbit_training(
model, 
use_gradient_checkpointing=gradient_checkpointing
)
# Prepare Peft model by adding Lora
peft_config = LoraConfig(
r=64,
lora_alpha=16,
target_modules=modules,
lora_dropout=0.1,
bias="none",
task_type=TaskType.CAUSAL_LM,
)
model = get_peft_model(model, peft_config)           

這樣模型就可以在本地的GPU上進行微調了。通過建立SFTTrainer (Trainer的一個子類,可以處理我們到目前為止讨論的所有内容)使這個過程變得更加容易。

from trl import SFTTrainer
model = AutoModelForCausalLM.from_pretrained(
"EleutherAI/gpt-neo-125m",
load_in_4bit=True,
device_map="auto",
)
trainer = SFTTrainer(
model,
train_dataset=dataset,
dataset_text_field="text",
torch_dtype=torch.bfloat16,
peft_config=peft_config,
)
trainer.train()           

總結

在本文中,介紹了大型語言模型微調過程中出現的一個挑戰:如何在單個GPU上進行微調。我們介紹了3種技術來減少記憶體占用:梯度檢查點、LoRA和量化。我們看到了如何通過利用PEFT、BitsAndBytes和Transformers将這些技術應用到我們的代碼中。

本文的目标是提供一個深入而簡單的視圖,利用的現有技術,以便在你的項目中微調自己的llm。

作者:Jeremy Arancio

繼續閱讀