深度學習模型訓練的效率與穩定性,仰賴學習率的調整和梯度裁剪等技術。學習率控制模型引數更新的幅度,梯度裁剪則避免梯度爆炸問題,兩者對於大語言模型(LLM)的訓練尤為關鍵。本文除了探討學習率暖啟動與餘弦衰減的原理及實作,也詳細說明梯度裁剪的計算方式與程式碼範例,更進一步介紹了低秩適應(LoRA)技術,這是一種引數高效的微調方法,可以減少儲存需求並提高可擴充套件性。文章最後也示範瞭如何準備資料集、載入預訓練模型、建立 PyTorch 資料集和資料載入器,為讀者提供實務操作的參考。
強化訓練迴圈的技術:學習率調整與梯度裁剪
在深度學習模型訓練中,特別是在大語言模型(LLM)的訓練過程中,學習率的調整和梯度裁剪是兩個至關重要的技術。本文將探討這兩種技術的原理、實作方法及其在模型訓練中的重要性。
學習率暖啟動與餘弦衰減
學習率暖啟動(Warmup)是一種在訓練初期逐漸增加學習率的技術,目的是避免模型在訓練初期因學習率過高或過低而導致的不穩定。接著,餘弦衰減(Cosine Decay)技術在暖啟動後逐漸降低學習率,以半個餘弦週期的軌跡下降至接近零。
程式碼實作:
import math
min_lr = 0.1 * initial_lr
track_lrs = []
lr_increment = (peak_lr - initial_lr) / warmup_steps
global_step = -1
for epoch in range(n_epochs):
for input_batch, target_batch in train_loader:
optimizer.zero_grad()
global_step += 1
if global_step < warmup_steps:
lr = initial_lr + global_step * lr_increment
else:
progress = ((global_step - warmup_steps) /
(total_training_steps - warmup_steps))
lr = min_lr + (peak_lr - min_lr) * 0.5 * (
1 + math.cos(math.pi * progress)
)
for param_group in optimizer.param_groups:
param_group["lr"] = lr
track_lrs.append(optimizer.param_groups[0]["lr"])
內容解密:
- 學習率初始化:根據
initial_lr和peak_lr設定初始學習率和峰值學習率。 - 暖啟動階段:在前
warmup_steps步中,學習率線性增加至peak_lr。 - 餘弦衰減階段:超出
warmup_steps後,學習率依照餘弦函式逐漸下降至min_lr。 - 學習率更新:每一步更新最佳化器的學習率,並記錄當前學習率。
梯度裁剪
梯度裁剪(Gradient Clipping)是一種透過限制梯度的範數(norm)來穩定模型訓練的技術。當梯度的範數超過設定的閾值時,將梯度按比例縮放至該閾值內。
L2 範數計算範例:
對於一個向量 ( v = [v_1, v_2, …, v_n] ),其 L2 範數定義為: [ ||v||_2 = \sqrt{v_1^2 + v_2^2 + … + v_n^2} ]
若有梯度矩陣: [ G = \begin{bmatrix} 3 & 4 \ 0 & 0 \end{bmatrix} ] 其 L2 範數為: [ ||G||_2 = \sqrt{3^2 + 4^2 + 0^2 + 0^2} = 5 ]
若設定的 max_norm 為 1,則縮放後的梯度矩陣 ( G’ ) 為:
[ G’ = G \times \frac{1}{5} = \begin{bmatrix} 0.6 & 0.8 \ 0 & 0 \end{bmatrix} ]
程式碼實作:
from chapter05 import calc_loss_batch
torch.manual_seed(123)
model = GPTModel(GPT_CONFIG_124M)
model.to(device)
loss = calc_loss_batch(input_batch, target_batch, model, device)
loss.backward()
def find_highest_gradient(model):
max_grad = None
for param in model.parameters():
if param.grad is not None:
grad_values = param.grad.data.flatten()
max_grad_param = grad_values.max()
if max_grad is None or max_grad_param > max_grad:
max_grad = max_grad_param
return max_grad
print(find_highest_gradient(model))
內容解密:
- 模型初始化:建立模型並將其移至指定裝置。
- 損失計算:計算一個批次的損失並進行反向傳播。
find_highest_gradient函式:遍歷模型引數,找出最大的梯度值。- 梯度檢查:輸出模型中最大的梯度值,用於檢查梯度是否過大。
D.4 修改後的訓練函式
前面的程式碼識別出的最大梯度值為 tensor(0.0411)
現在,讓我們套用梯度裁剪(gradient clipping)並觀察它如何影響最大的梯度值:
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
print(find_highest_gradient(model))
套用最大範數(max norm)為1的梯度裁剪後,最大的梯度值明顯變小: tensor(0.0185)
內容解密:
torch.nn.utils.clip_grad_norm_函式用於限制梯度的最大範數,防止梯度爆炸。max_norm=1.0設定了梯度的最大範數為1.0,超過此值的梯度將被縮放。- 梯度裁剪後的最大梯度值(0.0185)明顯小於之前的最大值(0.0411),顯示梯度裁剪的有效性。
修改後的訓練函式
最後,我們改進了 train_model_simple 訓練函式(參見第5章),加入了三個新概念:線性預熱(linear warmup)、餘弦衰減(cosine decay)和梯度裁剪。這些方法共同幫助穩定大語言模型(LLM)的訓練。
以下是修改後的程式碼,與 train_model_simple 的差異以註解標示:
from chapter05 import evaluate_model, generate_and_print_sample
def train_model(model, train_loader, val_loader, optimizer, device,
n_epochs, eval_freq, eval_iter, start_context, tokenizer,
warmup_steps, initial_lr=3e-05, min_lr=1e-6):
train_losses, val_losses, track_tokens_seen, track_lrs = [], [], [], []
tokens_seen, global_step = 0, -1
peak_lr = optimizer.param_groups[0]["lr"]
total_training_steps = len(train_loader) * n_epochs
lr_increment = (peak_lr - initial_lr) / warmup_steps
for epoch in range(n_epochs):
model.train()
for input_batch, target_batch in train_loader:
optimizer.zero_grad()
global_step += 1
if global_step < warmup_steps:
lr = initial_lr + global_step * lr_increment
else:
progress = ((global_step - warmup_steps) /
(total_training_steps - warmup_steps))
lr = min_lr + (peak_lr - min_lr) * 0.5 * (
1 + math.cos(math.pi * progress))
# 調整學習率
for param_group in optimizer.param_groups:
param_group["lr"] = lr
track_lrs.append(lr)
loss = calc_loss_batch(input_batch, target_batch, model, device)
loss.backward()
if global_step > warmup_steps:
# 套用梯度裁剪
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()
tokens_seen += input_batch.numel()
if global_step % eval_freq == 0:
train_loss, val_loss = evaluate_model(
model, train_loader, val_loader,
device, eval_iter
)
train_losses.append(train_loss)
val_losses.append(val_loss)
track_tokens_seen.append(tokens_seen)
print(f"Ep {epoch+1} (Iter {global_step:06d}): "
f"Train loss {train_loss:.3f}, "
f"Val loss {val_loss:.3f}")
generate_and_print_sample(
model, tokenizer, device, start_context
)
return train_losses, val_losses, track_tokens_seen, track_lrs
內容解密:
train_model函式結合了線性預熱、餘弦衰減和梯度裁剪,以穩定LLM的訓練。- 在預熱階段,學習率線性增加;之後,學習率按照餘弦衰減策略調整。
- 梯度裁剪在預熱階段後套用,以防止梯度爆炸。
- 訓練過程中,定期評估模型在訓練集和驗證集上的損失,並生成樣本文字。
訓練模型
定義 train_model 函式後,我們可以類別似於使用 train_model_simple 方法來訓練模型:
import tiktoken
torch.manual_seed(123)
model = GPTModel(GPT_CONFIG_124M)
model.to(device)
peak_lr = 5e-4
optimizer = torch.optim.AdamW(model.parameters(), weight_decay=0.1)
tokenizer = tiktoken.get_encoding("gpt2")
n_epochs = 15
train_losses, val_losses, tokens_seen, lrs = train_model(
model, train_loader, val_loader, optimizer, device, n_epochs=n_epochs,
eval_freq=5, eval_iter=1, start_context="Every effort moves you",
tokenizer=tokenizer, warmup_steps=warmup_steps,
initial_lr=1e-5, min_lr=1e-5
)
內容解密:
- 使用
train_model函式訓練模型,設定了相關的超引數,如n_epochs、eval_freq和warmup_steps。 - 訓練過程中,模型的訓練損失和驗證損失會被記錄和輸出。
- 鼓勵讀者使用更大的文字資料集進行訓練,並比較使用
train_model和train_model_simple的結果。
訓練結果
在MacBook Air或類別似筆記型電腦上,訓練大約需要5分鐘,並輸出以下結果:
Ep 1 (Iter 000000): Train loss 10.934, Val loss 10.939
Ep 1 (Iter 000005): Train loss 9.151, Val loss 9.461
Every effort moves you,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
...
Ep 15 (Iter 000130): Train loss 0.041, Val loss 6.915
Every effort moves you?" "Yes--quite insensible to the irony. She wanted him vindicated--and by me!" He laughed again, and threw back his head to look up at the sketch of the donkey. "There were days when I
與預訓練類別似,由於資料集較小且多次迭代,模型在幾個epoch後開始過擬合。儘管如此,我們可以看到函式正在工作,因為它最小化了訓練集損失。
低秩適應(LoRA)技術介紹
低秩適應(LoRA)是一種由 Hu 等人提出的引數高效微調方法(https://arxiv.org/abs/2106.09685)。LoRA 透過學習兩個遠小於原始權重矩陣 $W$ 的矩陣 $A$ 和 $B$ 來近似權重更新 $\Delta W$,其中 $\Delta W \approx AB$,$A$ 和 $B$ 之間進行矩陣乘法運算。
LoRA 的工作原理
使用 LoRA,我們可以將權重更新重新表述為: [ h = Wx + \Delta Wx = Wx + ABx ] 其中 $x$ 是輸入資料,$W$ 是預訓練權重,$\Delta W$ 是權重更新,$A$ 和 $B$ 是 LoRA 矩陣。
圖 E.1 直觀地展示了全微調和 LoRA 的權重更新公式之間的比較。
LoRA 的優勢
LoRA 的優勢在於它能夠在保持預訓練模型權重不變的情況下,透過動態應用 LoRA 矩陣來實作模型定製。這大大減少了儲存需求並提高了可擴充套件性,因為只需要調整和儲存較小的 LoRA 矩陣。
準備資料集
在將 LoRA 應用於垃圾郵件分類別示例之前,我們需要載入資料集和預訓練模型。以下是資料準備的程式碼:
from pathlib import Path
import pandas as pd
from ch06 import (
download_and_unzip_spam_data,
create_balanced_dataset,
random_split
)
url = "https://archive.ics.uci.edu/static/public/228/sms+spam+collection.zip"
zip_path = "sms_spam_collection.zip"
extracted_path = "sms_spam_collection"
data_file_path = Path(extracted_path) / "SMSSpamCollection.tsv"
download_and_unzip_spam_data(url, zip_path, extracted_path, data_file_path)
df = pd.read_csv(data_file_path, sep="\t", header=None, names=["Label", "Text"])
balanced_df = create_balanced_dataset(df)
balanced_df["Label"] = balanced_df["Label"].map({"ham": 0, "spam": 1})
train_df, validation_df, test_df = random_split(balanced_df, 0.7, 0.1)
train_df.to_csv("train.csv", index=None)
validation_df.to_csv("validation.csv", index=None)
test_df.to_csv("test.csv", index=None)
內容解密:
這段程式碼下載並解壓縮垃圾郵件資料集,然後建立平衡的資料集,並將其分為訓練集、驗證集和測試集,最後儲存為 CSV 檔案。
載入預訓練模型
接下來,我們載入預訓練的 GPT 模型:
from gpt_download import download_and_load_gpt2
from chapter04 import GPTModel
from chapter05 import load_weights_into_gpt
CHOOSE_MODEL = "gpt2-small (124M)"
INPUT_PROMPT = "Every effort moves"
BASE_CONFIG = {
"vocab_size": 50257,
"context_length": 1024,
"drop_rate": 0.0,
"qkv_bias": True
}
內容解密:
這段程式碼指定了要使用的預訓練 GPT 模型的組態,包括詞彙大小、上下文長度、丟棄率和查詢鍵值偏差等引數。
建立 PyTorch 資料集物件
import torch
from torch.utils.data import Dataset
import tiktoken
from chapter06 import SpamDataset
tokenizer = tiktoken.get_encoding("gpt2")
train_dataset = SpamDataset("train.csv", max_length=None, tokenizer=tokenizer)
val_dataset = SpamDataset("validation.csv", max_length=train_dataset.max_length, tokenizer=tokenizer)
test_dataset = SpamDataset("test.csv", max_length=train_dataset.max_length, tokenizer=tokenizer)
內容解密:
這段程式碼建立了 PyTorch 資料集物件,分別用於訓練、驗證和測試。SpamDataset 類別負責載入資料並進行必要的預處理。
建立 PyTorch 資料載入器
from torch.utils.data import DataLoader
num_workers = 0
batch_size = 8
torch.manual_seed(123)
train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers, drop_last=True)
val_loader = DataLoader(dataset=val_dataset, batch_size=batch_size, num_workers=num_workers, drop_last=False)
test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, num_workers=num_workers, drop_last=False)
內容解密:
這段程式碼建立了 PyTorch 資料載入器,用於批次載入資料。DataLoader 類別提供了靈活的資料載入方式,可以設定批次大小、工作執行緒數等引數。
驗證資料載入器
print("Train loader:")
for input_batch, target_batch in train_loader:
pass
print("Input batch dimensions:", input_batch.shape)
print("Label batch dimensions", target_batch.shape)
內容解密:
這段程式碼驗證了資料載入器的正確性,列印出了輸入批次和標籤批次的維度。
列印批次數量
print(f"{len(train_loader)} training batches")
print(f"{len(val_loader)} validation batches")
print(f"{len(test_loader)} test batches")
內容解密:
這段程式碼列印出了訓練、驗證和測試資料載入器中的批次數量。