從零開始實現大語言模型（十四）：高階訓練技巧

1. 前言

預訓練大語言模型的流程與訓練普通神經深度網絡模型本質上并沒有任何不同。可以使用深度學習實踐中已經被證明非常有效的高階訓練技巧，優化大語言模型預訓練流程，使大語言模型預訓練效率更高，訓練過程更穩定。

本文介紹深度學習領域優化訓練學習率的兩種方法Learning Rate Warmup和Cosine Decay，優化深度神經網絡模型參數梯度的方法Gradient Clipping，以及優化訓練超參數的方法Hyperparameters Search，并實現預訓練大語言模型的函數hyper_pretrain_model。

2. 優化訓練學習率

2.1 Learning Rate Warmup

學習率是訓練深度學習模型過程中最關鍵的超參數，沒有之一。學習率可以控制深度神經網絡模型參數的迭代更新速度，學習率越大，則參數的迭代更新速度越快，學習率越小，則參數的更新速度越慢。但是過大的學習率會導致損失函數在Error Surface上發生跳躍，使訓練過程不穩定，模型難以收斂。如果如學習率太小，則會導致參數深度神經網絡參數每次更新的幅度很小，使神經網絡模型的訓練效率很低，而且容易使損失函數陷入Error Surface中的局部最優解。

Learning Rate Warmup（學習率預熱）是一種經過深度學習實踐證明非常有效的優化深度神經網絡前幾次迭代訓練學習率，以降低深度神經網絡參數隨機初始化帶來的不確定性風險，從而提升訓練過程穩定性的方法。Learning Rate Warmup會指定一個非常小的初始學習率initial_lr，以及預熱步驟warmup_steps，并在訓練深度神經網絡的前warmup_steps次迭代流程中，將學習率逐步從initial_lr提升至不使用Learning Rate Warmup時設定的值peak_lr。在深度學習實踐中，預熱步驟warmup_steps一般會占總訓練次數的 $0.1\%$ 至 $10\%$ 。

如下面的代碼所示，假設學習率的設定值peak_lr為0.01，Learning Rate Warmup指定的初始學習率initial_lr為0.0001，warmup_steps為15。使用Learning Rate Warmup優化訓練學習率，需要在訓練深度神經網絡的前warmup_steps次迭代流程中，計算當次迭代訓練使用的學習率大小，并修改優化器中使用學習率的值：

import os
import torch
import random
import tiktoken
from torch.utils.data import Dataset, DataLoader# from [從零開始實現大語言模型（二）：文本數據處理] import LLMDataset
# from [從零開始實現大語言模型（七）：多頭注意力機制] import MultiHeadAttention
# from [從零開始實現大語言模型（八）：Layer Normalization] import LayerNorm
# from [從零開始實現大語言模型（九）：前饋神經網絡與GELU激活函數] import GELU, FeedForward
# from [從零開始實現大語言模型（十一）：構建大語言模型GPTModel] import TransformerBlock, GPTModeltorch.manual_seed(123)train_data_path = "train_data"
vocabulary = "gpt2"
special_token_id = 50256
context_len = 1024
stride = 1024
batch_size = 2embedding_dim = 768
num_layers = 12
num_heads = 12
context_len = 1024
vocabulary_size = 50257
dropout = 0.1
qkv_bias = Falsenum_epochs = 15
initial_lr = 0.0001
peak_lr = 0.01
warmup_steps = 15train_dataset = LLMDataset(train_data_path, vocabulary, special_token_id, context_len, stride)
train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True, drop_last=True)gpt2_small = GPTModel(embedding_dim=embedding_dim,num_layers=num_layers,num_heads=num_heads,context_len=context_len,vocabulary_size=vocabulary_size,dropout=dropout,qkv_bias=qkv_bias
)optimizer = torch.optim.AdamW(gpt2_small.parameters(), weight_decay=0.1)
lr_increment = (peak_lr - initial_lr) / warmup_stepsglobal_step = -1
track_lrs = []for epoch in range(num_epochs):for input_batch, target_batch in train_loader:optimizer.zero_grad()global_step += 1if global_step < warmup_steps:lr = initial_lr + global_step * lr_incrementelse:lr = peak_lrfor param_group in optimizer.param_groups:param_group["lr"] = lrtrack_lrs.append(optimizer.param_groups[0]["lr"])

可以使用如下代碼繪制大語言模型gpt2_small的每輪迭代訓練過程所使用的學習率：

import matplotlib.pyplot as pltplt.ylabel("Learning rate")
plt.xlabel("Step")
total_training_steps = len(train_loader) * num_epochs
plt.plot(range(total_training_steps), track_lrs)
plt.show()

執行上面代碼，生成大語言模型gpt2_small的整個迭代訓練流程中的學習率變化情況圖像如下：

2.2 Cosine Decay

在深度學習實踐中，一般會將Learning Rate Warmup與Cosine Decay結合起來，共同優化訓練學習率。Learning Rate Warmup只作用于深度神經網絡的前warmup_steps輪迭代訓練過程（預熱階段），使訓練學習率從一個很小的initial_lr逐步提升至peak_lr。Cosine Decay會在預熱階段之后的全部迭代訓練過程中，以余弦曲線的方式逐步減小訓練學習率，以降低模型參數的更新速度，減少損失函數越過Error Surface上極小值的概率，提升訓練過程穩定性。

如下面的代碼所示，在預熱階段之后使用Cosine Decay策略調整訓練學習率，需要使用global_step - warmup_steps得到預熱階段后的迭代訓練步數，以及使用total_training_steps - warmup_steps計算出預熱階段后的總迭代訓練次數，并通過(global_step - warmup_steps) / (total_training_steps - warmup_steps)計算出去掉預熱階段后的訓練進度百分比progress。使用math.cos(math.pi * progress)可以計算得到一個介于1到-1之間的余弦值，當progress為0時，余弦值為1，當progress為1時，余弦值為-1，余弦值的變化速率曲線為余弦曲線。使用0.5 * (1 + math.cos(math.pi * progress))對余弦值做變換，使余弦值的取值范圍由[1, -1]變換到[1, 0]，最后使用 min_lr + (peak_lr - min_lr) * 0.5 * (1 + math.cos(math.pi * progress))計算得到當前迭代訓練步數對應的訓練學習率：

import mathmin_lr = 0.1 * initial_lrtrack_lrs = []
global_step = -1for epoch in range(num_epochs):for input_batch, target_batch in train_loader:optimizer.zero_grad()global_step += 1if 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"] = lrtrack_lrs.append(optimizer.param_groups[0]["lr"])

可以使用如下代碼繪制使用Learning Rate Warmup與Cosine Decay策略后的學習率變化情況圖像：

plt.ylabel("Learning rate")
plt.xlabel("Step")
plt.plot(range(total_training_steps), track_lrs)
plt.show()

執行上面代碼，生成的學習率變化情況圖像如下：

3. 優化模型參數梯度

3.1 Gradient Clipping

Gradient Clipping（梯度裁剪）是一種通過限制參數梯度大小，以解決深度神經網絡訓練過程中的梯度爆炸問題，從而提升訓練過程穩定性的模型參數梯度優化方法。深度學習實踐中常用的Gradient Clipping方法有兩種：基于梯度值的裁剪和基于梯度范數的裁剪。

基于梯度值的裁剪方法的原理非常簡單，其會直接將深度神經網絡參數的梯度中大于clip_value的梯度設置成clip_value，并將小于-clip_value的梯度設置成-clip_value，使深度神經網絡參數梯度的絕對值值不超過clip_value。基于梯度范數的裁剪方法首先會計算神經網絡參數梯度的p-范數，如果p-范數大于max_norm，則會將每個梯度值均乘以 $max_norm p_norm \frac{\text{max\_norm}}{\text{p\_norm}}$ ，使神經網絡參數梯度的p-范數等于max_norm。

假設深度神經網絡共包含4個參數，后向傳播流程計算出的參數梯度 $G=\begin{bmatrix}1&2\\2&4\end{bmatrix}$ ，使用基于梯度范數的裁剪方法優化模型參數梯度，設置參數梯度2-范數的最大值max_norm為2.0，首先需要計算神經網絡參數梯度的2-范數 $\|G\|_2=\sqrt{1^2+2^2+2^2+4^2}=\sqrt{25}=5$ 。因為 $G\|_2>2$ ，因此會將每個梯度值均乘以 $max_norm p_norm = 2 5 \frac{\text{max\_norm}}{\text{p\_norm}}=\frac{2}{5}$ ，即將神經網絡參數梯度裁剪成 $G'=\frac{2}{5}\times G=\begin{bmatrix}\frac{2}{5}&\frac{4}{5}\\ \frac{4}{5}&\frac{8}{5}\end{bmatrix}$ 。

如下面的代碼所示，定義計算深度神經網絡參數梯度最大值的函數find_largest_gradient，并使用torch.tensor函數創建訓練樣本input_batch及訓練樣本標簽target_batch。將訓練樣本input_batch輸入大語言模型gpt2_small，使用calc_loss_batch函數計算大語言模型的預測輸出與訓練樣本標簽之間的交叉熵損失loss，并通過loss.backward()計算大語言模型參數梯度。最后使用find_largest_gradient函數打印輸入大語言模型參數梯度的最大值：

# from [從零開始實現大語言模型（十三）：預訓練大語言模型GPTModel] import calc_loss_batchdef find_largest_gradient(model):max_grad = Nonefor 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_paramreturn max_graddevice = torch.device("cpu")
input_batch = torch.tensor([[16833, 3626, 6100],  # [["every effort moves"],[40, 1107, 588]]      # ["I really like"]]
)
target_batch = torch.tensor([[3626, 6100, 345],  # [[" effort moves you"],[588, 428, 11311]]  # [" really like chocolate"]]
)loss = calc_loss_batch(input_batch, target_batch, gpt2_small, device)
loss.backward()
print(find_largest_gradient(gpt2_small))

執行上面代碼，打印結果如下：

tensor(0.6413)

使用上述基于梯度范數的裁剪方法優化模型參數梯度，設置大語言模型參數梯度的2-范數最大值max_norm為1.0，并打印經過Gradient Clipping優化之后的大語言模型參數梯度的最大值：

torch.nn.utils.clip_grad_norm_(gpt2_small.parameters(), max_norm=1.0)print(find_largest_gradient(gpt2_small))

執行上面代碼，打印結果如下：

tensor(0.0348)

4. 實現高階預訓練函數

可以結合上述3種高階訓練技巧實現預訓練大語言模型的函數hyper_pretrain_model。修改前文從零開始實現大語言模型（十三）：預訓練大語言模型GPTModel中實現的預訓練大語言模型的函數pretrain_model，在每輪for循環使用calc_loss_batch函數計算大語言模型的預測輸出與訓練樣本標簽之間的交叉熵損失之前，先使用2中所述優化訓練學習率的兩種方法Learning Rate Warmup和Cosine Decay，計算當次迭代訓練使用的學習率大小，并修改訓練優化器中使用學習率的值。在使用optimizer.step()方法更新大語言模型參數之前，先使用3中所述優化模型參數梯度的方法Gradient Clipping，優化模型參數梯度。具體代碼如下所示：

# from [從零開始實現大語言模型（十三）：預訓練大語言模型GPTModel] import calc_loss_loader, val_and_savedef hyper_pretrain_model(model, optimizer, train_loader, num_epochs, device, eval_freq, eval_iter, tokenizer, start_context, save_freq, checkpoint_dir, warmup_steps=10, initial_lr=3e-05, min_lr=1e-6, max_norm=1.0,checkpoint=None, val_loader=None
):if not os.path.exists(checkpoint_dir):os.makedirs(checkpoint_dir, exist_ok=True)if checkpoint is not None:model_checkpoint_path = os.path.join(checkpoint_dir, f"model_{checkpoint:06d}.pth")optimizer_checkpoint_path = os.path.join(checkpoint_dir, f"optimizer_{checkpoint:06d}.pth")model.load_state_dict(torch.load(model_checkpoint_path))optimizer.load_state_dict(torch.load(optimizer_checkpoint_path))else:checkpoint = -1train_losses, val_losses, track_tokens_seen, track_lrs = [], [], [], []tokens_seen, global_step = 0, -1peak_lr = optimizer.param_groups[0]["lr"]total_training_steps = len(train_loader) * num_epochslr_increment = (peak_lr - initial_lr) / warmup_stepsfor epoch in range(num_epochs):model.train()for i, (input_batch, target_batch) in enumerate(train_loader):if global_step % eval_freq == 0:model.train()optimizer.zero_grad()global_step += 1if 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"] = lrtrack_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=max_norm)optimizer.step()tokens_seen += input_batch.numel()print(f"Epoch {epoch + 1} (Batch {i:06d}): Train loss {loss.item():.3f}")checkpoint, train_loss, val_loss = val_and_save(model, optimizer, train_loader, val_loader, epoch, global_step, eval_freq,eval_iter, start_context, tokenizer, save_freq, checkpoint_dir, checkpoint, device)if train_loss is not None:train_losses.append(train_loss)val_losses.append(val_loss)track_tokens_seen.append(tokens_seen)checkpoint, _, _ = val_and_save(model, optimizer, train_loader, val_loader, epoch, global_step, 1,eval_iter, start_context, tokenizer, 1, checkpoint_dir, checkpoint, device)print(f"Epoch {epoch + 1} finished, checkpoint: {checkpoint:06d}")return train_losses, val_losses, track_tokens_seen, track_lrs

5. 優化訓練超參數

5.1 Hyper-parameters Search

超參數(hyper-parameters)是指需要在搭建和訓練深度神經網絡之前手動設置的一些參數。在深度學習中，有兩類超參數，一類超參數是深度神經網絡結構超參數，比如深度神經網絡的層數，Embedding向量的維度等等。另一類超參數是訓練超參數，例如訓練深度神經網絡使用的學習率，每個batch中訓練樣本數量等等。

優化深度神經網絡結構超參數的方法被統稱為神經網絡結構搜索(NAS, neural architecture search)。神經網絡結構搜索方法大致可分類3類，其中一類被稱為“大海撈針”，即根據實踐經驗定義一個有限的超參數搜索空間，逐一使用搜索空間中的超參數組合構建并訓練深度神經網絡直至收斂，取驗證集上測試指標最高的超參數組合作為搜索結果。另一類是不可微方法，其一般會將驗證集上的測試指標作為環境給的獎勵，使用強化學習算法搜索出較優的超參數組合。還有一類是可微方法，其核心思想是定義一個神經網絡結構超參數的可微函數作為目標函數，基于Super-net對目標函數關于超參數求梯度，直接使用梯度更新超參數。

本文不會詳細介紹深度神經網絡結構超參數優化方法，不同大語言模型的結構基本相同，Embedding向量維度等結構超參數一般會取決于可用的計算資源，工業界實踐中一般不會使用神經網絡架構搜索方法確定大語言模型的結構超參數。《從零開始實現大語言模型》系列專欄全部完成之后，我應該會寫幾篇博客詳細神經網絡結構搜索，感興趣的讀者可以關注我的個人博客。

預訓練大語言模型的時間成本及計算成本都非常高，例如訓練大語言模型Llama 2的數據共包含2T（萬億）個tokens，花費184320 A100 GPU時，換算成云計算資源價值，大約需要690000美元。在預訓練大語言模型的工業界實踐中，一般會在正式開始訓預訓練大語言模型之前，在相對小的數據集上，使用Hyper-parameters Search得到一個比較好的訓練超參數組合。Hyper-parameters Search的核心思想就是“大海撈針”，即定義一個有限的超參數搜索空間HPARAM_GRID，逐一使用搜索空間中的超參數組合訓練大語言模型，取驗證集上交叉熵損失最小的超參數組合作為正式預訓練大語言模型時所用的訓練超參數。具體代碼如下所示：

import itertoolsdef hparams_search_train(model, optimizer, train_loader, val_loader, num_epochs, device,eval_iter, warmup_steps, initial_lr, min_lr, max_norm
):global_step = -1peak_lr = optimizer.param_groups[0]["lr"]total_training_steps = len(train_loader) * num_epochslr_increment = (peak_lr - initial_lr) / warmup_stepsfor epoch in range(num_epochs):model.train()for input_batch, target_batch in train_loader:optimizer.zero_grad()global_step += 1if global_step < warmup_steps:lr = initial_lr + global_step * lr_incrementelse: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"] = lrloss = 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=max_norm)optimizer.step()train_loss = calc_loss_loader(train_loader, model, device, eval_iter)val_loss = calc_loss_loader(val_loader, model, device, eval_iter)return train_loss, val_lossHPARAM_GRID = {"batch_size": [2, 4, 8, 16],"dropout": [0.0, 0.1, 0.2],"warmup_steps": [10, 20, 30],"weight_decay": [0.1, 0.01, 0.0],"max_norm": [1.0, 0.5, 2.0],"peak_lr": [0.0001, 0.0005, 0.001, 0.005],"initial_lr": [0.00005, 0.0001],"min_lr": [0.00005, 0.00001, 0.0001],"num_epochs": [5, 10, 15, 20, 25],
}
hyperparameter_combinations = list(itertools.product(*HPARAM_GRID.values()))
print(f"Total hyperparameter configurations: {len(hyperparameter_combinations)}")device = torch.device("cpu")
val_data_path = "val_data"
train_dataset = LLMDataset(train_data_path, vocabulary, special_token_id, context_len, stride)
val_dataset = LLMDataset(val_data_path, vocabulary, special_token_id, context_len, stride)best_val_loss, best_train_loss = float("inf"), float("inf")
best_hparams = {}for i, combination in enumerate(hyperparameter_combinations):print(f"Evaluating configuration {i + 1} of {len(hyperparameter_combinations)}")HPARAM_CONFIG = dict(zip(HPARAM_GRID.keys(), combination))torch.manual_seed(123)train_loader = DataLoader(dataset=train_dataset, batch_size=HPARAM_CONFIG["batch_size"], shuffle=True, drop_last=True)val_loader = DataLoader(dataset=val_dataset, batch_size=HPARAM_CONFIG["batch_size"], shuffle=False, drop_last=False)model = GPTModel(embedding_dim=embedding_dim, num_layers=num_layers, num_heads=num_heads, context_len=context_len,vocabulary_size=vocabulary_size, dropout=HPARAM_CONFIG["dropout"], qkv_bias=qkv_bias)model.to(device)optimizer = torch.optim.AdamW(model.parameters(), lr=HPARAM_CONFIG["peak_lr"],weight_decay=HPARAM_CONFIG["weight_decay"])train_loss, val_loss = hparams_search_train(model, optimizer, train_loader, val_loader, HPARAM_CONFIG["num_epochs"], device, eval_iter=1,warmup_steps=HPARAM_CONFIG["warmup_steps"], initial_lr=HPARAM_CONFIG["initial_lr"],min_lr=HPARAM_CONFIG["min_lr"], max_norm=HPARAM_CONFIG["max_norm"])if val_loss < best_val_loss:best_val_loss = val_lossbest_train_loss = train_lossbest_hparams = HPARAM_CONFIGprint(f"Evaluating configuration {i + 1} completed.")print(f"Current best hyper-parameters: {best_hparams}")print(f"Current best Val loss: {best_val_loss} | Training loss {best_train_loss}")print("============================================================================")print("Hyper-parameter search completed.")
print(f"Best hyper-parameters: {best_hparams}")
print(f"Best Val loss: {best_val_loss} | Training loss {best_train_loss}")

神經網絡結構搜索領域的不可微方法并不適用于大語言模型訓練超參數搜索，訓練大語言模型所需的計算量太大，且使用強化學習算法搜索超參數，需要從頭開始完整訓練一次大語言模型才能獲得1個獎勵，強化學習算法一般至少需要上萬至數十萬次獎勵反饋才能收斂。

神經網絡結構搜索領域的可微方法同樣不適用于大語言模型訓練超參數搜索，所有可微方法的核心思想都是定義一個神經網絡結構超參數的可微函數作為目標函數，然而基本沒有辦法找到一個神經網絡訓練超參數的可微函數。