从零开始实现大语言模型（十四）：高阶训练技巧

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，则会将每个梯度值均乘以$\frac{\text{max\_norm}}{\text{p\_norm}}$，使神经网络参数梯度的p-范数等于max_norm。

假设深度神经网络共包含4个参数，后向传播流程计算出的参数梯度$G=\left[\begin{array}{cc}1& 2\\ 2& 4\end{array}\right]$，使用基于梯度范数的裁剪方法优化模型参数梯度，设置参数梯度2-范数的最大值max_norm为2.0，首先需要计算神经网络参数梯度的2-范数$\Vert G{\Vert }_2=\sqrt{1^2+2^2+2^2+4^2}=\sqrt{25}=5$。因为$\Vert G{\Vert }_2>2$，因此会将每个梯度值均乘以$\frac{\text{max\_norm}}{\text{p\_norm}}=\frac{2}{5}$，即将神经网络参数梯度裁剪成$G^{\prime }=\frac{2}{5}\times G=\left[\begin{array}{cc}\frac{2}{5}& \frac{4}{5}\\ \frac{4}{5}& \frac{8}{5}\end{array}\right]$。

如下面的代码所示，定义计算深度神经网络参数梯度最大值的函数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个奖励，强化学习算法一般至少需要上万至数十万次奖励反馈才能收敛。

神经网络结构搜索领域的可微方法同样不适用于大语言模型训练超参数搜索，所有可微方法的核心思想都是定义一个神经网络结构超参数的可微函数作为目标函数，然而基本没有办法找到一个神经网络训练超参数的可微函数。

6. 结束语

预训练大语言模型的流程与训练普通神经深度网络模型本质上并没有任何不同，其难度不在于算法，而在于数据，更在于算力。绝大部分企业都没有预训练大语言模型的算力资源，因此如何利用开源大语言模型成了大语言模型工业实践中的重中之重，接下来一起看看如何加载开源大语言模型参数吧！