从零开始实现大语言模型（十三）：预训练大语言模型GPTModel

1. 前言

使用梯度下降算法通过下一个token预测任务预训练大语言模型GPTModel，前向传播流程每次会输入一个batch的长度均为context_len的训练样本，执行 $batch_size × context_len \text{batch\_size}\times\text{context\_len}$ 次下一个token预测任务，共预测输出 $batch_size × context_len \text{batch\_size}\times\text{context\_len}$ 个tokens。后向传播流程首先会使用交叉熵(Cross Entropy)损失函数计算大语言模型GPTModel的预测输出与训练样本标签之间的损失(loss)，再通过后向传播算法计算大语言模型参数梯度，最后使用梯度下降算法更新大语言模型的参数。

本文使用交叉熵损失函数计算生成大语言模型GPTModel的预测输出与训练样本标签之间的loss，介绍大语言模型的预训练流程，并实现预训练大语言模型的函数pretrain_model。

2. 损失函数Cross Entropy

交叉熵损失函数可以度量大语言模型的预测输出与训练样本标签之间的差异。损失函数计算生成的loss值越大，表明大语言模型的预测输出与训练样本标签之间的差异越大，loss值越小，表明大语言模型的预测输出与训练样本标签之间的差异越小。

对输入文本做tokenization，将输入文本转换成包含context_len个token ID的列表，并输入大语言模型GPTModel，可以得到context_len个维度为vocabulary_size的logits向量，第 $i$ 个logits向量是大语言模型根据前 $i$ 个token预测生成的下一个token的概率分数向量，logits向量中的第 $k$ 个概率分数值越大，表明大语言模型预测生成的下一个token的ID为 $k$ 的概率越高。

使用softmax函数将大语言模型预测生成的logits向量归一化，得到大语言模型预测生成的下一个token的概率分布，概率分布中对应样本标签位置的概率值表示大语言模型预测输出的token为相应训练样本标签的概率。对应样本标签位置的概率值越接近1，表明大语言模型预测输出的token为相应训练样本标签的概率越高，大语言模型的预测输出与训练样本标签之间的差异越小。

使用梯度下降算法预训练大语言模型GPTModel的前向传播流程中，大语言模型每次会预测生成 $batch_size × context_len \text{batch\_size}\times\text{context\_len}$ 个下一个token的概率分布。如下图所示，交叉熵损失函数会分别获取 $batch_size × context_len \text{batch\_size}\times\text{context\_len}$ 个概率分布中对应样本标签位置的概率值，使用对数函数计算这些概率值的对数，并计算所有对数值的均值，最后将对数均值的相反数作为大语言模型GPTModel的预测输出与训练样本标签之间的损失loss。

如下面的代码所示，使用torch.tensor函数创建训练样本inputs及训练样本标签targets，将训练样本inputs输入大语言模型gpt2_small，并使用softmax函数将大语言模型的输出张量logits归一化，得到 $2\times3$ 个下一个token的概率分布，其中每个概率分布的维度均等于词汇表的大小50257。分别获取 $2\times3$ 个下一个token的概率分布中对应样本标签位置的概率值，使用torch.log函数计算这些概率值的对数，并计算所有对数值均值的相反数，可以得到大语言模型gpt2_small的预测输出与样本标签targets之间的交叉熵损失：

import torch
# from [从零开始实现大语言模型（七）：多头注意力机制] import MultiHeadAttention
# from [从零开始实现大语言模型（八）：Layer Normalization] import LayerNorm
# from [从零开始实现大语言模型（九）：前馈神经网络与GELU激活函数] import GELU, FeedForward
# from [从零开始实现大语言模型（十一）：构建大语言模型GPTModel] import TransformerBlock, GPTModel

torch.manual_seed(123)

embedding_dim = 768
num_layers = 12
num_heads = 12
context_len = 1024
vocabulary_size = 50257
dropout = 0.1
qkv_bias = False

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
)

inputs = torch.tensor(
    [[16833, 3626, 6100],  # [["every effort moves"],
     [40, 1107, 588]]      # ["I really like"]]
)

targets = torch.tensor(
    [[3626, 6100, 345],  # [[" effort moves you"],
     [588, 428, 11311]]  # [" really like chocolate"]]
)

with torch.no_grad():
    logits = gpt2_small(inputs)
probas = torch.softmax(logits, dim=-1)

target_probas_1 = probas[0, [0, 1, 2], targets[0]]
target_probas_2 = probas[1, [0, 1, 2], targets[1]]

log_probas = torch.log(torch.cat((target_probas_1, target_probas_2)))
avg_log_probas = torch.mean(log_probas)
neg_avg_log_probas = avg_log_probas * -1

print("probas shape:", probas.shape)
print("target_probas_1:", target_probas_1)
print("target_probas_2:", target_probas_2)
print("log_probas:", log_probas)
print("avg_log_probas:", avg_log_probas)
print("cross entropy loss:", neg_avg_log_probas)

执行上面代码，打印结果如下：

probas shape: torch.Size([2, 3, 50257])
target_probas_1: tensor([2.6369e-05, 1.5997e-05, 1.6926e-05])
target_probas_2: tensor([1.5638e-05, 8.9422e-06, 1.7967e-05])
log_probas: tensor([-10.5433, -11.0431, -10.9867, -11.0658, -11.6247, -10.9270])
avg_log_probas: tensor(-11.0318)
cross entropy loss: tensor(11.0318)

可以直接使用PyTorch内置的cross_entropy函数执行上述计算流程，得到大语言模型gpt2_small的预测输出logits与样本标签targets之间的交叉熵损失：

loss = torch.nn.functional.cross_entropy(logits.flatten(0, 1), targets.flatten())
print(loss)

执行上面代码，打印结果如下：

tensor(11.0318)

根据打印结果可知，上述6个步骤计算得到的损失值与PyTorch内置的cross_entropy函数计算得到的交叉熵损失完全相同。交叉熵损失本质上就是大语言模型预测生成的 $batch_size × context_len \text{batch\_size}\times\text{context\_len}$ 个下一个token的概率分布中对应样本标签位置概率值的对数均值的相反数。

3. 大语言模型预训练流程

预训练大语言模型的流程与训练普通神经网络模型本质上并没有任何不同。如下图所示，预训练大语言模型可以把整个训练数据集扫几个epoch，每个epoch会把整个训练数据集扫一遍，每次会使用训练数据集中一个batch的训练样本训练一次大语言模型。前向传播流程会将一个batch的训练样本输入大语言模型，得到大语言模型的预测输出logits。后向传播流程首先会使用交叉熵损失函数计算大语言模型的预测输出logits与训练样本标签targets之间的损失loss，再通过后向传播算法计算大语言模型参数梯度，最后使用梯度下降算法更新大语言模型的参数。

可以使用如下代码定义计算一个batch样本数据交叉熵损失的函数calc_loss_batch，以及计算整个数据集上所有样本数据交叉熵损失的函数calc_loss_loader。函数calc_loss_batch将一个batch的样本数据输入大语言模型，得到大语言模型的预测输出logits，并使用torch.nn.functional.cross_entropy函数计算大语言模型的预测输出logits与样本标签targets之间的损失loss。函数calc_loss_loader每次取数据集中一个batch的样本数据，使用calc_loss_batch函数计算该batch样本数据的交叉熵损失，并返回数据集上所有样本数据损失loss的均值：

def calc_loss_batch(input_batch, target_batch, model, device):
    input_batch = input_batch.to(device)
    target_batch = target_batch.to(device)
    logits = model(input_batch)
    loss = torch.nn.functional.cross_entropy(
        logits.flatten(0, 1), target_batch.flatten()
    )
    return loss

def calc_loss_loader(data_loader, model, device, num_batches=None):
    model.eval()
    total_loss = 0.0
    if num_batches is None:
        num_batches = len(data_loader)
    else:
        num_batches = min(num_batches, len(data_loader))

    with torch.no_grad():
        for i, (input_batch, target_batch) in enumerate(data_loader):
            if i < num_batches:
                loss = calc_loss_batch(input_batch, target_batch, model, device)
                total_loss += loss.item()
            else:
                break
    return total_loss / num_batches

实现预训练大语言模型的函数pretrain_model，可以使用for循环将整个训练数据集扫num_epochs遍，并在每次训练大语言模型的循环中，首先使用optimizer.zero_grad函数将大语言模型所有参数的梯度置为0，然后使用函数calc_loss_batch计算一个batch训练样本的交叉熵损失loss。使用loss.backward函数可以执行后向传播流程，计算大语言模型所有参数的梯度，并通过optimizer.step函数使用梯度下降算法更新大语言模型参数。具体代码如下所示：

# from [从零开始实现大语言模型（十二）：文本生成策略] import generate_text

def pretrain_model(
        model, optimizer, train_loader, num_epochs, device,
        eval_freq, eval_iter, tokenizer, start_context,
        save_freq, checkpoint_dir, 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 = -1

    train_losses, val_losses, track_tokens_seen = [], [], []
    tokens_seen, global_step = 0, -1

    for 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()
            loss = calc_loss_batch(input_batch, target_batch, model, device)
            loss.backward()
            optimizer.step()

            tokens_seen += input_batch.numel()
            global_step += 1
            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


def 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
):
    train_loss, val_loss = None, None
    if global_step % eval_freq == 0:
        if val_loader is not None:
            train_loss = calc_loss_loader(train_loader, model, device, eval_iter)
            val_loss = calc_loss_loader(val_loader, model, device, eval_iter)
            print(f"Epoch {epoch + 1} (Step {global_step:06d}): Train loss {train_loss:.3f}, Val loss {val_loss:.3f}")
        generated_sample_text = generate_text(
            model, start_context, max_new_tokens=50, tokenizer=tokenizer,
            context_size=model.pos_emb.weight.shape[0], top_k=1, compact_format=True
        )
        print(f"Generated Sample Text: {generated_sample_text}")
        print("=====================================================================")

    if global_step % save_freq == 0:
        checkpoint += 1
        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")
        torch.save(model.state_dict(), model_checkpoint_path)
        torch.save(optimizer.state_dict(), optimizer_checkpoint_path)
    return checkpoint, train_loss, val_loss

PyTorch中神经网络模型的state_dict是一个字典对象，字典中的key为神经网络模型中参数的名称，value为相应的参数。使用.state_dict函数可以一次性获取神经网络模型中的所有参数，并通过torch.save函数将所有参数保存为一个checkpoint。torch.load函数可以读取指定checkpoint，通过.load_state_dict函数可以将神经网络模型中的参数修改为checkpoint中的记录值。

所有具有自适应能力的优化器（如AdamW可以根据历史梯度信息动态调整学习率）都需要记录每个神经网络参数的历史梯度等信息，同样可以使用.state_dict一次性获取优化器中的所有数据记录，以及通过.load_state_dict函数从指定checkpoint中还原这些记录数据。

如下面的代码所示，使用从零开始实现大语言模型（二）：文本数据处理中构建的Dataset创建训练集train_dataset及验证集val_dataset，并通过PyTorch内置的torch.utils.data.DataLoader类创建训练集及验证集对应的DataLoader。使用torch.optim.AdamW实例化训练大语言模型的优化器optimizer，最后使用函数pretrain_model预训练大语言模型gpt2_small：

import os
import random
import tiktoken
from torch.utils.data import Dataset, DataLoader

# from [从零开始实现大语言模型（二）：文本数据处理] import LLMDataset

train_data_path = "train_data"
val_data_path = "val_data"
vocabulary = "gpt2"
special_token_id = 50256
context_len = 1024
stride = 1024
batch_size = 2

num_epochs = 10
eval_freq = 5
eval_iter = 1
save_freq = 5
checkpoint_dir = "checkpoint"
start_context = "萧炎，斗之力，三段"
tokenizer = tiktoken.encoding_for_model(vocabulary)
device = torch.device("cpu")
gpt2_small.to(device)
optimizer = torch.optim.AdamW(gpt2_small.parameters(), lr=0.0006, weight_decay=0.1)

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)
train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True, drop_last=True)
val_loader = DataLoader(dataset=val_dataset, batch_size=batch_size, shuffle=False, drop_last=False)
print(f"train_loader len: {len(train_loader)}")

train_losses, val_losses, tokens_seen = pretrain_model(
    gpt2_small, optimizer, train_loader, num_epochs, device,
    eval_freq, eval_iter, tokenizer, start_context,
    save_freq, checkpoint_dir, val_loader=val_loader
)

执行上面代码，打印结果如下：

train_loader len: 7
Epoch 1 (Batch 000000): Train loss 11.034
Epoch 1 (Step 000000): Train loss 9.827, Val loss 9.784
Generated Sample Text: 萧炎，斗之力，三段 Knowledge�缌�缌缌�703 clashes�缌 longest，，缌，���，缌缌�
=====================================================================
Epoch 1 (Batch 000001): Train loss 9.940
Epoch 1 (Batch 000002): Train loss 8.811
Epoch 1 (Batch 000003): Train loss 7.954
Epoch 1 (Batch 000004): Train loss 7.286
Epoch 1 (Batch 000005): Train loss 6.629
Epoch 1 (Step 000005): Train loss 5.980, Val loss 6.003
Generated Sample Text: 萧炎，斗之力，三段，，，，，，，，，，，，，，，，�
=====================================================================
Epoch 1 (Batch 000006): Train loss 6.027
Epoch 1 (Step 000006): Train loss 5.390, Val loss 5.479
Generated Sample Text: 萧炎，斗之力，三段，，，�，，，��，，，，，�，�，，，
=====================================================================
Epoch 1 finished, checkpoint: 000002
Epoch 2 (Batch 000000): Train loss 5.401
Epoch 2 (Batch 000001): Train loss 5.028
Epoch 2 (Batch 000002): Train loss 4.788
Epoch 2 (Batch 000003): Train loss 4.616
Epoch 2 (Step 000010): Train loss 4.511, Val loss 4.526
Generated Sample Text: 萧炎，斗之力，三段，�，�，�，�，�，�，�，�，�，��，�，
=====================================================================

[...]

Epoch 9 (Step 000060): Train loss 2.561, Val loss 3.470
Generated Sample Text: 萧炎，斗之力，三段���是在脸�的�，�炣�殸废是萧炣也是曰�，萧�
=====================================================================
Epoch 9 (Batch 000005): Train loss 2.560
Epoch 9 (Batch 000006): Train loss 2.558
Epoch 9 (Step 000062): Train loss 2.456, Val loss 3.455
Generated Sample Text: 萧炎，斗之力，三段���，脸庿，炎�，萧炎萧�炎�萧�，萧�的�
=====================================================================
Epoch 9 finished, checkpoint: 000021
Epoch 10 (Batch 000000): Train loss 2.525
Epoch 10 (Batch 000001): Train loss 2.388
Epoch 10 (Batch 000002): Train loss 2.663
Epoch 10 (Step 000065): Train loss 2.270, Val loss 3.468
Generated Sample Text: 萧炎，斗之力，三段��技萧�的萧炣也�，萧�讵��更中着曰萧�着�
=====================================================================
Epoch 10 (Batch 000003): Train loss 2.464
Epoch 10 (Batch 000004): Train loss 2.602
Epoch 10 (Batch 000005): Train loss 2.511
Epoch 10 (Batch 000006): Train loss 2.557
Epoch 10 (Step 000069): Train loss 2.117, Val loss 3.474
Generated Sample Text: 萧炎，斗之力，三段��，这的�法的萧�炼�萧�法，萧�级级父了�
=====================================================================
Epoch 10 finished, checkpoint: 000023

从上面的打印结果可知，使用梯度下降算法训练大语言模型gpt2_small，可以减小大语言模型的预测输出与样本标签之间的交叉熵损失，并显著提升大语言模型的文本生成能力。在训练刚开始时，将萧炎，斗之力，三段输入大语言模型gpt2_small，生成的是Knowledge�缌�缌缌�703 clashes�缌 longest，，缌，��，缌缌�或者，，，，，，，，，，，，，，，，�这样不包含任何有效信息的自然语言文本序列。在仅包含7个batch训练样本的数据集上训练10个epoch，大语言模型gpt2_small已经可以生成��，脸庿，炎�，萧炎萧�炎�萧�，萧�的�以及��，这的�法的萧�炼�萧�法，萧�级级父了�这样与训练数据集存在一定关联的自然语言文本了。

可以使用如下代码，分别绘制大语言模型gpt2_small在训练集及验证集上交叉熵损失的变化情况图像：

import matplotlib.pyplot as plt

def plot_losses(epochs_seen, tokens_seen, train_losses, val_losses):
    fig, ax1 = plt.subplots(figsize=(5, 3))
    ax1.plot(epochs_seen, train_losses, label="Training loss")
    ax1.plot(epochs_seen, val_losses, linestyle="-.", label="Validation loss")
    ax1.set_xlabel("Epochs")
    ax1.set_ylabel("Loss")
    ax1.legend(loc="upper right")
    ax2 = ax1.twiny()
    ax2.plot(tokens_seen, train_losses, alpha=0)
    ax2.set_xlabel("Tokens seen")
    fig.tight_layout()
    plt.show()

epochs_tensor = torch.linspace(0, num_epochs, len(train_losses))
plot_losses(epochs_tensor, tokens_seen, train_losses, val_losses)

执行上面代码，生成交叉熵损失的变化情况图像如下：

从上面的交叉熵损失变化情况图像可知，在训练刚开始时，训练集及验证集上的交叉熵损失都非常大。使用梯度下降算法训练大语言模型gpt2_small，可以减小大语言模型的预测输出与样本标签之间的交叉熵损失，使大语言模型的预测输出与样本标签之间的差异性更小。

随着训练的进行，训练集和验证集上交叉熵损失的差异会越来越大，训练集上的交叉熵损失值会比验证集小的越来越明显，表明大语言模型在训练数据集上的过拟合情况越来越严重。在工业界的预训练大语言模型实践中，并不会在一个很小的训练数据集上训练多个epoch，而是会在一个非常大的训练数据集上训练少数几个甚至只训练一个epoch，这种训练策略可以很大程度上解决预训练大语言模型时的过拟合问题。

4. 结束语

前向传播流程将一个batch的训练样本输入大语言模型，共预测输出 $batch_size × context_len \text{batch\_size}\times\text{context\_len}$ 个维度为vocabulary_size的logits向量。后向传播流程首先使用交叉熵损失函数计算大语言模型的预测输出与训练样本标签之间的损失loss，并通过后向传播算法计算大语言模型参数梯度，最后使用梯度下降算法更新大语言模型的参数。

预训练大语言模型就是不断从训练数据集中获取一个batch的训练样本，然后执行这个操作直至收敛的过程。