【阅读记录-章节5】Build a Large Language Model (From Scratch)

目录

• 5. Pretraining on unlabeled data
• • 5.1 Evaluating generative text models
  • • 5.1.1 Evaluating generative text models
    • 5.1.2 Calculating the text generation loss
    • • 评估模型生成文本的质量
    • 5.1.3 Calculating the training and validation set losses
  • 5.2 Training an LLM
  • 5.3 Decoding strategies to control randomness
  • • 5.3.1 Temperature Scaling
    • 5.3.2 Top-k sampling
    • 5.3.3 Modifying the text generation function
    • • 练习5.2 补充介绍
      • 练习5.3 补充介绍
  • 5.4 Loading and saving model weights in PyTorch
  • • • 练习5.4 补充介绍
  • 5.5 Loading pretrained weights from OpenAI

---

5. Pretraining on unlabeled data

到目前为止，我们已经实现了数据采样和注意力机制，并编写了大型语言模型（LLM）的架构代码。接下来，我们将实现一个训练函数并对LLM进行预训练。同时，我们将学习基本的模型评估技术，以衡量生成文本的质量，这是在训练过程中优化LLM所必需的。此外，我们还将讨论如何加载预训练的权重，为我们的LLM在微调阶段提供一个坚实的起点。图5.1展示了我们的整体计划，重点介绍了本章将讨论的内容。

权重参数的补充介绍

• 在大型语言模型（LLM）和其他深度学习模型的上下文中，权重参数是模型学习和调整的核心。这些权重也被称为权重参数，或简称为参数，指的是可训练的参数，学习过程会对其进行调整，以优化模型的性能。

权重参数的作用

• 决定模型表现：权重参数决定了模型在输入数据上的表现，通过训练过程不断调整这些参数，以最小化预测误差，从而提升模型的准确性和生成文本的质量。

在PyTorch中的存储与访问

• 存储位置：在像PyTorch这样的深度学习框架中，权重通常存储在线性层（torch.nn.Linear）中。例如，我们在第3章中使用线性层实现了多头注意力模块，在第4章中实现了GPT模型。
• 访问方式：
  • 单个层的权重：初始化一个层（例如 new_layer = torch.nn.Linear(...)）后，可以通过 .weight 属性访问其权重，如 new_layer.weight。
  • 所有可训练参数：为了方便起见，PyTorch 允许通过 model.parameters() 方法直接访问模型的所有可训练参数，包括权重和偏置。这在实现训练循环时尤为重要，因为我们需要迭代这些参数以更新模型。

5.1 Evaluating generative text models

在本章中，我们将在简要回顾第4章的文本生成内容后，设置我们的大型语言模型（LLM）进行文本生成，并讨论评估生成文本质量的基本方法。随后，我们将计算训练损失和验证损失。图5.2展示了本章涵盖的主题，前三个步骤已被突出显示。

5.1.1 Evaluating generative text models

首先，让我们设置LLM并简要回顾我们在第4章中实现的文本生成过程。我们从初始化GPT模型开始，稍后将使用GPTModel类和GPT_CONFIG_124M字典对其进行评估和训练（参见第4章）：

import torch
from chapter04 import GPTModel

GPT_CONFIG_124M = {
    "vocab_size": 50257,
    "context_length": 256,
    "emb_dim": 768,
    "n_heads": 12,
    "n_layers": 12,
    "drop_rate": 0.1,
    "qkv_bias": False
}

torch.manual_seed(123)
model = GPTModel(GPT_CONFIG_124M)
model.eval()

我们将上下文长度（context_length）从1024个令牌缩短至256个令牌。相比上一章，唯一的调整是减少了上下文长度，这一修改降低了训练模型的计算需求，使得在标准笔记本电脑上进行训练成为可能。

原始的GPT-2模型拥有1.24亿个参数，配置为处理最多1024个令牌。在训练过程完成后，我们将更新上下文长度设置，并加载预训练权重，以使模型能够处理配置为1024个令牌上下文长度的模型。使用GPTModel实例，我们采用第4章中的generate_text_simple函数，并引入两个便捷函数：text_to_token_ids和token_ids_to_text。这些函数有助于在文本和令牌表示之间进行转换，这是我们将在本章中贯穿使用的技术。

图5.3展示了使用GPT模型的三步文本生成过程：

1. 令牌化：分词器将输入文本转换为一系列令牌ID（参见第2章）。
2. 模型生成：模型接收这些令牌ID并生成相应的logits，这些logits是表示词汇表中每个令牌概率分布的向量（参见第4章）。
3. 解码：这些logits被转换回令牌ID，分词器将其解码为人类可读的文本，完成从文本输入到文本输出的循环。

我们可以按照以下代码实现文本生成过程：

Listing 5.1 文本与令牌ID转换的实用函数

import tiktoken
from chapter04 import generate_text_simple

def text_to_token_ids(text, tokenizer):
    encoded = tokenizer.encode(text, allowed_special={'<|endoftext|>'})
    # 移除批次维度
    encoded_tensor = torch.tensor(encoded).unsqueeze(0)
    return encoded_tensor

def token_ids_to_text(token_ids, tokenizer):
    flat = token_ids.squeeze(0)
    return tokenizer.decode(flat.tolist())

start_context = "Every effort moves you"
tokenizer = tiktoken.get_encoding("gpt2")
token_ids = generate_text_simple(
    model=model,
    idx=text_to_token_ids(start_context, tokenizer),
    max_new_tokens=10,
    context_size=GPT_CONFIG_124M["context_length"]
)
print("Output text:\n", token_ids_to_text(token_ids, tokenizer))

使用这段代码，模型生成了以下文本：

Output text:
Every effort moves you rentingetic wasn? refres RexMeCHicular stren

显然，模型尚未生成连贯的文本，因为它尚未经过训练。为了定义什么使得文本“连贯”或“高质量”，我们必须实现一种数值方法来评估生成的内容。这种方法将使我们能够在整个训练过程中监控和提升模型的性能。

接下来，我们将计算生成输出的损失指标。这个损失将作为训练进展的指标。此外，在后续章节中，当我们对LLM进行微调时，我们将回顾评估模型质量的其他方法。

5.1.2 Calculating the text generation loss

在训练大型语言模型（LLM）时，评估生成文本的质量是至关重要的。通过计算文本生成的损失，我们可以量化模型的性能，并指导训练过程的改进。本文将通过一个实际的例子，逐步讲解如何评估模型生成文本的质量。

我们将从回顾数据的加载方式和 generate_text_simple 函数生成文本的过程开始。

如图5.4所示，文本生成过程可以分为五个步骤：

1. 输入文本转换为Token IDs：将输入文本转换为对应的Token IDs序列。
2. 模型预测下一个Token的概率分布：将Token IDs输入模型，计算每个位置下一个Token的概率分布（Logits），并通过Softmax函数转换为概率。
3. 选择最可能的下一个Token：对概率分布应用Argmax函数，选取概率最高的Token ID。
4. 更新Token序列：将选取的Token ID添加到输入序列中，重复步骤2和3，生成完整的输出序列。
5. 将Token IDs转换回文本：将生成的Token IDs序列转换回可读的文本。

需要注意的是，图5.4中的示例为了简化，仅使用了一个包含7个Token的小词汇表。然而，在实际中，我们的GPT模型使用了包含50,257个词的更大词汇表，因此Token IDs的范围是0到50,256。

为了更好地理解，我们使用两个输入示例：

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

对应的目标输出（Targets）为：

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

注意，目标输出是将输入序列右移一个位置得到的。这种方式在第2章的数据加载器实现中已介绍过，主要用于训练模型预测序列中的下一个Token。

我们将输入数据输入模型，计算Logits，并通过Softmax函数转换为概率：

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

输出的概率张量形状为：

torch.Size([2, 3, 50257])

• 第一个维度2表示批次大小，即有两个输入示例。
• 第二个维度3表示每个输入序列中的Token数量。
• 第三个维度50257表示词汇表的大小，即每个位置上可能的下一个Token的概率分布。

通过对概率分布应用Argmax函数，我们可以得到模型预测的下一个Token IDs：

token_ids = torch.argmax(probas, dim=-1, keepdim=True)
print("Token IDs:\n", token_ids)

输出的Token IDs为：

Token IDs:
tensor([[[16657],
         [  339],
         [42826]],

        [[49906],
         [29669],
         [41751]]])

我们使用Tokenizer将预测的Token IDs转换回可读文本：

print(f"Targets batch 1: {token_ids_to_text(targets[0], tokenizer)}")
print(f"Outputs batch 1: {token_ids_to_text(token_ids[0].flatten(), tokenizer)}")

输出结果：

Targets batch 1: effort moves you
Outputs batch 1: Armed heNetflix

可以看到，模型生成的文本与目标文本有明显差异，这是因为模型还没有经过训练，输出的结果较为随机。

评估模型生成文本的质量

现在，我们希望通过计算损失（如图5.5所示）来数值化地评估模型生成文本的性能。这不仅有助于衡量生成文本的质量，也是实现训练函数的基础，我们将使用它来更新模型的权重，以改进生成的文本。

我们实现的文本评估过程的一部分，如图5.5所示，是测量生成的Tokens与正确预测（目标）之间的“距离”。我们稍后实现的训练函数将使用这些信息来调整模型的权重，使其生成的文本更接近（或理想情况下匹配）目标文本。

模型训练的目标是提高与正确目标Token IDs对应的索引位置上的Softmax概率，如图5.6所示。我们接下来将实现的评估指标也使用了这个Softmax概率来数值化地评估模型生成的输出：在正确位置上的概率越高，效果越好。

请记住，图5.6显示的是针对一个包含7个Token的小词汇表的Softmax概率，以便将所有内容放在一张图中。这意味着初始的随机值大约在1/7左右，约等于0.14。然而，我们的GPT-2模型使用的词汇表有50,257个Token，因此大多数初始概率将徘徊在0.00002（1/50,257）左右。

为了量化模型的性能，我们需要计算生成文本与目标文本之间的差异。这通常通过计算损失函数（如交叉熵损失）来实现。

我们可以提取模型在目标Token位置上的概率值：

# 设置文本索引为0，表示第一个文本样本
text_idx = 0

# 提取第一个文本样本中每个位置上目标Token的概率
# probas的形状为[批次大小, 序列长度, 词汇表大小]
# targets[text_idx]包含了第一个文本样本的目标Token IDs
# [0, 1, 2]表示序列中的三个位置
target_probas_1 = probas[text_idx, [0, 1, 2], targets[text_idx]]

# 打印第一个文本样本的目标Token概率
print("Text 1:", target_probas_1)

# 设置文本索引为1，表示第二个文本样本
text_idx = 1

# 提取第二个文本样本中每个位置上目标Token的概率
# targets[text_idx]包含了第二个文本样本的目标Token IDs
target_probas_2 = probas[text_idx, [0, 1, 2], targets[text_idx]]

# 打印第二个文本样本的目标Token概率
print("Text 2:", target_probas_2)

输出结果：

Text 1: tensor([7.4541e-05, 3.1061e-05, 1.1563e-05])
Text 2: tensor([1.0337e-05, 5.6776e-05, 4.7559e-06])

这些概率值表示模型在每个位置上预测目标Token的概率。由于模型尚未训练，这些概率值非常低。

训练模型的目标是最大化正确Token的概率，即提高目标Token在概率分布中的概率值。这可以通过最小化损失函数来实现。

常用的损失函数是交叉熵损失，它可以衡量预测分布与目标分布之间的差异。具体的计算方法将在后续的训练函数中实现。

反向传播

• 我们如何最大化与目标 Token 对应的 Softmax 概率值？总体而言，我们需要更新模型的权重，使模型对我们希望生成的相应 Token ID 输出更高的概率值。权重的更新是通过一种称为反向传播的过程完成的，这是训练深度神经网络的标准技术（有关反向传播和模型训练的更多详细信息，请参见附录 A 的 A.3 至 A.7 节）。
• 反向传播需要一个损失函数，该函数计算模型预测输出（在这里是对应目标 Token ID 的概率值）与实际期望输出之间的差异。这个损失函数衡量了模型的预测与目标值之间的偏差有多大。

接下来，我们将为两个示例批次 target_probas_1 和 target_probas_2 计算损失。主要步骤如图5.7所示。

计算概率分数的对数

由于我们已经应用了步骤1到3以获得 target_probas_1 和 target_probas_2，接下来进行步骤4，对概率分数取对数：

# 对目标概率取对数并合并两个批次

# 1. 使用 torch.cat 函数将两个目标概率张量 (target_probas_1 和 target_probas_2) 在第一个维度（即批次维度）上进行拼接
#    - target_probas_1 和 target_probas_2 分别对应两个不同的输入批次
#    - 拼接后的张量形状为 [6]，因为每个批次有3个概率值，总共2个批次
combined_probas = torch.cat((target_probas_1, target_probas_2))  

# 2. 对拼接后的概率值取自然对数
#    - torch.log 函数计算每个元素的自然对数
#    - 对数转换有助于将乘法操作转化为加法，简化后续的数学优化过程
#    - 取对数后，较小的概率值会变得更负，便于计算损失
log_probas = torch.log(combined_probas)  

# 3. 打印对数概率值以供检查
print(log_probas) 

输出结果为：

tensor([ -9.5042, -10.3796, -11.3677, -11.4798, -9.7764, -12.2561])

在数学优化中，处理概率分数的对数比直接处理概率分数更为便捷。虽然本书不深入探讨这一主题，但在附录B的讲座中有更详细的介绍。

接下来，我们将这些对数概率合并为一个单一的分数，通过计算平均值（图5.7的步骤5）：

# 计算对数概率的平均值
avg_log_probas = torch.mean(log_probas)
print(avg_log_probas)

输出结果为：

tensor(-10.7940)

我们的目标是通过更新模型的权重，使平均对数概率尽可能接近0。然而，在深度学习中，常见的做法不是将平均对数概率提升到0，而是将负的平均对数概率降低到0。负的平均对数概率就是将平均对数概率乘以-1，对应图5.7的步骤6：

# 计算负的平均对数概率
neg_avg_log_probas = avg_log_probas * -1
print(neg_avg_log_probas)

输出结果为：

tensor(10.7940)

在深度学习中，将这个负值（-10.7940）转换为正值（10.7940）的过程称为交叉熵损失。幸运的是，PyTorch 已经内置了 cross_entropy 函数，可以为我们完成图5.7中的所有六个步骤。

交叉熵损失

• 交叉熵损失是机器学习和深度学习中常用的一种度量方法，用于衡量两个概率分布之间的差异——通常是真实标签的分布（这里是数据集中的Token）与模型预测的分布（例如，LLM生成的Token概率）。在机器学习的框架中，特别是像PyTorch这样的框架，cross_entropy 函数计算离散结果的交叉熵，这类似于模型生成的Token概率下目标Token的负平均对数概率，使得“交叉熵”和“负平均对数概率”这两个术语在实际中相关且常常可以互换使用。

在应用 cross_entropy 函数之前，让我们简要回顾一下 Logits 和 Targets 张量的形状：

print("Logits shape:", logits.shape)
print("Targets shape:", targets.shape)

输出结果为：

Logits shape: torch.Size([2, 3, 50257])
Targets shape: torch.Size([2, 3])

可以看出，Logits 张量有三个维度：批次大小、Token 数量和词汇表大小。Targets 张量有两个维度：批次大小和 Token 数量。

对于 PyTorch 中的 cross_entropy 损失函数，我们需要通过合并批次维度来将这些张量展平：

# 将Logits和Targets展平以适应cross_entropy函数
logits_flat = logits.flatten(0, 1)
targets_flat = targets.flatten()
print("Flattened logits:", logits_flat.shape)
print("Flattened targets:", targets_flat.shape)

输出结果为：

Flattened logits: torch.Size([6, 50257])
Flattened targets: torch.Size([6])

请记住，Targets 是我们希望 LLM 生成的 Token IDs，Logits 包含了在进入 Softmax 函数之前模型未缩放的输出。

之前，我们应用了 Softmax 函数，选择了对应目标 ID 的概率分数，并计算了负的平均对数概率。现在，PyTorch 的 cross_entropy 函数将为我们完成所有这些步骤：

# 计算交叉熵损失
loss = torch.nn.functional.cross_entropy(logits_flat, targets_flat)
print(loss)

输出结果为：

tensor(10.7940)

这个损失值与我们手动应用图5.7中的各个步骤时得到的结果相同。

困惑度（Perplexity）

• 困惑度是常与交叉熵损失一起使用的度量，用于评估模型在语言建模等任务中的性能。它提供了一种更易于理解的方式来理解模型在预测序列中下一个 Token 时的不确定性。
• 困惑度衡量模型预测的概率分布与数据集中实际单词分布的匹配程度。与损失类似，较低的困惑度表明模型的预测更接近实际分布。
• 困惑度通常被认为比原始的损失值更具可解释性，因为它表示模型在每一步预测下一个 Token 时不确定的有效词汇量。在给定的示例中，这意味着模型在词汇表中的48,725个Token中对生成下一个 Token 的选择感到不确定。

我们已经为两个小型文本输入计算了损失，用于说明目的。接下来，我们将把损失计算应用到整个训练集和验证集上，以进一步评估和优化模型的性能。

5.1.3 Calculating the training and validation set losses

在训练大型语言模型（LLM）之前，我们必须首先准备训练和验证数据集。接下来，如图5.8所示，我们将计算训练集和验证集的交叉熵损失，这是模型训练过程中一个重要的组成部分。

为了计算训练和验证数据集的损失，我们使用一个非常小的文本数据集——埃迪丝·沃顿（Edith Wharton）的短篇小说《裁决》（The Verdict），这是我们在第2章中已经使用过的。选择公共领域的文本可以避免任何使用权相关的问题。此外，使用这样一个小的数据集可以让代码示例在标准笔记本电脑上几分钟内执行完成，即使没有高端GPU，这对于教育目的尤其有利。

注意：感兴趣的读者还可以使用本书的补充代码准备一个由古腾堡计划（Project Gutenberg）中超过60,000本公共领域书籍组成的大规模数据集，并在其上训练LLM（详细信息见附录D）。

以下代码加载《裁决》短篇小说：

# 加载《裁决》短篇小说
import os
import urllib.request

if not os.path.exists("the-verdict.txt"):
    url = ("https://raw.githubusercontent.com/rasbt/"
           "LLMs-from-scratch/main/ch02/01_main-chapter-code/"
           "the-verdict.txt")
    file_path = "the-verdict.txt"
    urllib.request.urlretrieve(url, file_path)
    
with open(file_path, "r", encoding="utf-8") as file:
    text_data = file.read()

加载数据集后，我们可以检查数据集中的字符数和Token数：

# 计算数据集中的字符数和Token数
total_characters = len(text_data)
total_tokens = len(tokenizer.encode(text_data))
print("Characters:", total_characters)
print("Tokens:", total_tokens)

输出结果为：

Characters: 20479
Tokens: 5145

尽管只有5,145个Token，文本可能看起来太小，不足以训练一个LLM，但如前所述，这是为了教育目的，以便我们可以在几分钟内运行代码而不是数周。此外，稍后我们将从OpenAI加载预训练权重到我们的GPTModel代码中。

接下来，我们将数据集划分为训练集和验证集，并使用第2章的数据加载器准备LLM训练的批次。这个过程在图5.9中进行了可视化。

由于空间限制，我们使用了max_length=6。然而，对于实际的数据加载器，我们将max_length设置为LLM支持的256 Token上下文长度，以便在训练过程中让LLM看到更长的文本。

注意：为了简化和提高效率，我们使用了大小相似的数据块进行训练。然而，实际中，使用可变长度输入训练LLM也有助于模型在使用时更好地泛化不同类型的输入。

首先，我们定义一个train_ratio，使用90%的数据进行训练，剩下的10%作为验证数据用于模型评估：

# 定义训练集比例
train_ratio = 0.90

# 计算划分索引
split_idx = int(train_ratio * len(text_data))

# 划分训练集和验证集
train_data = text_data[:split_idx]
val_data = text_data[split_idx:]

使用train_data和val_data子集，我们现在可以创建各自的数据加载器，重用第2章中的create_dataloader_v1代码：

# 从第2章导入create_dataloader_v1函数
from chapter02 import create_dataloader_v1

# 设置随机种子以确保可重复性
torch.manual_seed(123)

# 创建训练数据加载器
train_loader = create_dataloader_v1(
    train_data,
    batch_size=2,
    max_length=GPT_CONFIG_124M["context_length"],
    stride=GPT_CONFIG_124M["context_length"],
    drop_last=True,
    shuffle=True,
    num_workers=0
)

# 创建验证数据加载器
val_loader = create_dataloader_v1(
    val_data,
    batch_size=2,
    max_length=GPT_CONFIG_124M["context_length"],
    stride=GPT_CONFIG_124M["context_length"],
    drop_last=False,
    shuffle=False,
    num_workers=0
)

我们使用了相对较小的批次大小来减少计算资源需求，因为我们使用的是一个非常小的数据集。实际上，训练LLM时使用1,024或更大的批次大小并不罕见。

作为一个可选检查步骤，我们可以迭代数据加载器以确保它们被正确创建：

# 打印训练数据加载器的形状
print("Train loader:")
for x, y in train_loader:
    print(x.shape, y.shape)

# 打印验证数据加载器的形状
print("\nValidation loader:")
for x, y in val_loader:
    print(x.shape, y.shape)

我们应该看到如下输出：

Train loader:
torch.Size([2, 256]) torch.Size([2, 256])
torch.Size([2, 256]) torch.Size([2, 256])
torch.Size([2, 256]) torch.Size([2, 256])
torch.Size([2, 256]) torch.Size([2, 256])
torch.Size([2, 256]) torch.Size([2, 256])
torch.Size([2, 256]) torch.Size([2, 256])
torch.Size([2, 256]) torch.Size([2, 256])
torch.Size([2, 256]) torch.Size([2, 256])
torch.Size([2, 256]) torch.Size([2, 256])

Validation loader:
torch.Size([2, 256]) torch.Size([2, 256])

基于上述代码输出，我们有九个训练集批次，每个批次包含两个样本和256个Token。由于我们仅分配了10%的数据用于验证，因此只有一个验证批次，包含两个输入示例。如预期，输入数据（x）和目标数据（y）具有相同的形状（批次大小乘以每批次的Token数），因为目标是将输入序列右移一个位置得到的，如第2章所述。

接下来，我们实现一个实用函数，用于计算通过训练和验证加载器返回的给定批次的交叉熵损失：

def calc_loss_batch(input_batch, target_batch, model, device):
    """
    计算单个批次的交叉熵损失

    参数：
    - input_batch (Tensor): 输入批次的Token IDs
    - target_batch (Tensor): 目标批次的Token IDs
    - model (nn.Module): 训练中的LLM模型
    - device (torch.device): 计算设备（CPU或GPU）

    返回：
    - loss (Tensor): 该批次的交叉熵损失
    """
    # 将输入和目标批次移动到指定设备
    input_batch = input_batch.to(device)
    target_batch = target_batch.to(device)
    
    # 获取模型的Logits输出
    logits = model(input_batch)
    
    # 计算交叉熵损失
    loss = torch.nn.functional.cross_entropy(
        logits.flatten(0, 1),  # 展平Logits以适应cross_entropy函数
        target_batch.flatten()  # 展平目标批次
    )
    
    return loss

我们可以使用calc_loss_batch实用函数来实现以下calc_loss_loader函数，该函数计算通过给定数据加载器采样的所有批次的损失：

def calc_loss_loader(data_loader, model, device, num_batches=None):
    """
    计算通过数据加载器（data_loader）返回的多个批次的平均交叉熵损失。

    参数：
    - data_loader (DataLoader): 数据加载器，提供输入和目标批次。
    - model (nn.Module): 要评估的模型。
    - device (torch.device): 计算设备（CPU或GPU）。
    - num_batches (int, 可选): 要处理的批次数。如果为None，则处理所有批次。

    返回：
    - avg_loss (float): 所有处理批次的平均交叉熵损失。
    """

    # 初始化总损失为0
    total_loss = 0.0

    # 如果数据加载器为空，返回NaN
    if len(data_loader) == 0:
        return float("nan")
    
    # 如果未指定处理的批次数，则处理所有批次
    elif num_batches is None:
        num_batches = len(data_loader)
    
    # 如果指定了批次数，则取较小的值以避免超出数据加载器的范围
    else:
        num_batches = min(num_batches, len(data_loader))
    
    # 迭代数据加载器中的批次
    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
    
    # 计算并返回平均损失
    avg_loss = total_loss / num_batches if num_batches > 0 else float("nan")
    return avg_loss

现在，我们已经准备好了计算训练和验证集损失的工具函数，接下来我们将整个流程整合起来，计算训练和验证集的交叉熵损失。

# 指定计算设备（GPU 或 CPU）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 将模型移动到指定设备
model.to(device)

# 禁用梯度计算，因为我们只需要前向传播来计算损失
with torch.no_grad():
    # 计算训练集的平均交叉熵损失
    train_loss = calc_loss_loader(train_loader, model, device)
    
    # 计算验证集的平均交叉熵损失
    val_loss = calc_loss_loader(val_loader, model, device)
    
# 打印训练集和验证集的损失
print("Training loss:", train_loss)
print("Validation loss:", val_loss)

输出示例：

Training Loss: 10.7940
Validation Loss: 10.7940

这些损失值反映了模型在训练集和验证集上的表现。较低的损失值表示模型的预测与目标更接近。

现在我们有了一种衡量生成文本质量的方法，我们将训练LLM以减少这一损失，从而使其在生成文本方面表现得更好，如图5.10所示。

接下来，我们将专注于对LLM进行预训练。模型训练完成后，我们还将实现替代的文本生成策略，并学习如何保存和加载预训练模型的权重。

5.2 Training an LLM

终于到了实现预训练LLM（我们的GPTModel）代码的时候。为此，我们将专注于一个简单明了的训练循环，以保持代码简洁易读。

注意：感兴趣的读者可以在附录D中了解更多高级技术，包括学习率预热（learning rate warmup）、余弦退火（cosine annealing）和梯度裁剪（gradient clipping）。

图5.11中的流程图展示了一个典型的PyTorch神经网络训练工作流程，我们将使用它来训练LLM。该流程图概述了八个步骤，从每个epoch的迭代、处理批次、重置梯度、计算损失和新梯度、更新权重，到最后的监控步骤，如打印损失和生成文本样本。

注意：如果您对使用PyTorch训练深度神经网络相对陌生，并且对这些步骤中的任何一个不熟悉，建议阅读附录A的A.5至A.8节。

我们可以通过以下代码中的 train_model_simple 函数来实现这个训练流程。

def train_model_simple(model, train_loader, val_loader,
                       optimizer, device, num_epochs,
                       eval_freq, eval_iter, start_context, tokenizer):
    """
    计算通过数据加载器（data_loader）返回的多个批次的平均交叉熵损失。
    
    参数：
    - model (nn.Module): 要评估的模型。
    - train_loader (DataLoader): 训练数据加载器。
    - val_loader (DataLoader): 验证数据加载器。
    - optimizer (Optimizer): 优化器，用于更新模型权重。
    - device (torch.device): 计算设备（CPU或GPU）。
    - num_epochs (int): 训练的总轮数。
    - eval_freq (int): 评估的频率（每多少步进行一次评估）。
    - eval_iter (int): 在每次评估中处理的批次数。
    - start_context (str): 用于生成文本样本的起始文本片段。
    - tokenizer (Tokenizer): 分词器，用于编码和解码文本。
    
    返回：
    - train_losses (list): 训练集的损失记录。
    - val_losses (list): 验证集的损失记录。
    - track_tokens_seen (list): 已见Token数量的记录。
    """
    # 初始化用于跟踪损失和已见Token数量的列表
    train_losses, val_losses, track_tokens_seen = [], [], []
    tokens_seen, global_step = 0, -1
    
    # 遍历每个epoch
    for epoch in range(num_epochs):
        model.train()  # 设置模型为训练模式