大模型原理:训练大语言模型

训练流程

上图8个步骤,从遍历每个训练轮次开始,处理批次,重置梯度,计算损失和新梯度,更新权重,最后以监控步骤(包括打印损失、生成文本样本等操作)结束。

实现

准备代码

需要 大模型原理:计算训练集和验证集的损失 文中的代码

定义训练函数

def train_model_simple(model, train_loader, val_loader, optimizer, device, num_epochs,
                       eval_freq, eval_iter, start_context, tokenizer):
    # Initialize lists to track losses and tokens seen
    train_losses, val_losses, track_tokens_seen = [], [], []
    tokens_seen, global_step = 0, -1

    # Main training loop
    for epoch in range(num_epochs):
        model.train()  # Set model to training mode

        for input_batch, target_batch in train_loader:
            optimizer.zero_grad() # Reset loss gradients from previous batch iteration
            loss = calc_loss_batch(input_batch, target_batch, model, device)
            loss.backward() # Calculate loss gradients
            optimizer.step() # Update model weights using loss gradients
            tokens_seen += input_batch.numel()
            global_step += 1

            # Optional evaluation step
            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} (Step {global_step:06d}): "
                      f"Train loss {train_loss:.3f}, Val loss {val_loss:.3f}")

        # Print a sample text after each epoch
        generate_and_print_sample(
            model, tokenizer, device, start_context
        )

    return train_losses, val_losses, track_tokens_seen

evaluate_model()函数

该函数对应于上图中的第(7)步。它会在每次模型更新后打印训练集和验证集的损失,以便可以 评估 训练是否改善了模型性能。

注意:由于该函数是 评估模型性能,所以要开启 .eval()( 禁用 Dropout)和 no_grad()(禁用梯度跟踪)

def evaluate_model(model, train_loader, val_loader, device, eval_iter):
    # 在评估阶段禁用dropout,以生成稳定且可复现的结果
    model.eval()
    # 评估阶段也会禁用梯度跟踪,因为这是不需要的,而且这样可以减少计算开销
    with torch.no_grad():
        train_loss = calc_loss_loader(train_loader, model, device, num_batches=eval_iter)
        val_loss = calc_loss_loader(val_loader, model, device, num_batches=eval_iter)
    model.train()
    return train_loss, val_loss

generate_and_print_sample()函数

该函数,用来跟踪模型在训练过程中是否有所改进。
主要功能:该函数以文本片段(start_context)作为输入,先将其转换为 词元ID,然后将其提供给大语言模型,使用之前定义的 generate_text_simple() 函数生成一个文本样本(此时是 词元ID),最后转为文本

注意:由于该函数是 评估模型性能,所以要开启 .eval()( 禁用 Dropout)和 no_grad()(禁用梯度跟踪)

def generate_and_print_sample(model, tokenizer, device, start_context):
    # 在评估阶段禁用dropout,以生成稳定且可复现的结果
    model.eval()
    context_size = model.pos_emb.weight.shape[0]
    encoded = text_to_token_ids(start_context, tokenizer).to(device)
    # 评估阶段也会禁用梯度跟踪,因为这是不需要的,而且这样可以减少计算开销
    with torch.no_grad():
        token_ids = generate_text_simple(
            model=model, idx=encoded,
            max_new_tokens=50, context_size=context_size
        )
        decoded_text = token_ids_to_text(token_ids, tokenizer)
        print(decoded_text.replace("\n", " "))  # Compact print format
    model.train()

AdamW

Adam优化器是训练深度神经网络的一种常见选择。然而,我们的训练循环中选择了AdamW优化器。AdamW是Adam的一个变体,它改进了权重衰减方法,旨在通过对较大的权重进行惩罚来最小化模型复杂性并防止过拟合。这种调整使得AdamW能够实现更有效的正则化和更好的泛化能力。因此,在大语言模型的训练中经常使用AdamW。

训练

对 GPTModel 实例进行 10 轮的训练:

torch.manual_seed(123)
model = GPTModel(GPT_CONFIG_124M)
model.to(device)
# ·parameters()方法返回模型的所有可训练权重参数
optimizer = torch.optim.AdamW(model.parameters(), lr=0.0004, weight_decay=0.1)

num_epochs = 10
train_losses, val_losses, tokens_seen = train_model_simple(
    model, train_loader, val_loader, optimizer, device,
    num_epochs=num_epochs, eval_freq=5, eval_iter=5,
    start_context="Every effort moves you", tokenizer=tokenizer
)

执行 train_model_simple() 函数会启动训练过程,在此执行过程中打印的输出如下所示:

训练集损失有了显著的改善,从 9.781 的初始值收敛到了 0.391。模型的语言能力得到了相当大的提升。在开始阶段,模型只能在起始上下文后添加逗号(Every effort moves you,,,,,,,,,,,,)或重复单词 and。在训练结束时,它已经可以生成语法正确的文本。

与训练集损失类似,验证集损失在训练过程中从较高值(9.933)开始逐渐降低。然而,它永远不会像训练集损失那样变得很小,在第 10 轮之后其值为 6.452。

图表展示损失

用图表展示训练集和验证集的损失

import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator


def plot_losses(epochs_seen, tokens_seen, train_losses, val_losses):
    fig, ax1 = plt.subplots(figsize=(5, 3))

    # Plot training and validation loss against epochs
    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")
    ax1.xaxis.set_major_locator(MaxNLocator(integer=True))  # only show integer labels on x-axis

    # Create a second x-axis for tokens seen
    ax2 = ax1.twiny()  # Create a second x-axis that shares the same y-axis
    ax2.plot(tokens_seen, train_losses, alpha=0)  # Invisible plot for aligning ticks
    ax2.set_xlabel("Tokens seen")

    fig.tight_layout()  # Adjust layout to make room
    plt.savefig("loss-plot.pdf")
    plt.show()

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

生成的训练集和验证集的损失图如下图所示。可以看到,训练集损失和验证集损失在第一轮开始改善。然而,损失在第二轮后开始发散。这种发散以及验证集损失远大于训练集损失的事实表明模型对训练数据过拟合。可以通过搜索生成的文本片段(比如 “The Verdict” 文本文件中的 quite insensible to the irony )来确认模型逐字记住了训练数据。

在训练开始阶段,训练集损失和验证集损失急剧下降,这表明模型正在学习。然而,在第二轮之后,训练集损失继续下降,验证集损失则停滞不前。这表明模型仍在学习,但在第二轮之后开始对训练集过拟合

这种记忆现象其实是可以预料到的,因为我们使用了一个非常非常小的训练数据集,并且对模型进行了多轮训练。通常,在更大的数据集上训练模型时,只训练一轮是很常见的做法。

总结

完成本章的前4项目标。接下来,我们将探讨用于大语言模型的文本生成策略,以减少训练数据的记忆,并增加大语言模型生成文本的独创性。在此之后,我们将讨论权重加载、保存以及从OpenAI的GPT模型加载预训练权重。

模型在完成训练后已经可以生成连贯的文本了


原文出处:/show_1GW2gao7GXBd.html