用 pytorch 从零开始创建大语言模型（五）：预训练无标注数据

5 预训练无标注数据

本章内容包括：

• 计算训练集和验证集的损失(loss)，用于评估 LLM 在训练过程中生成文本的质量。
• 实现训练函数(training function) 并预训练(pretraining) LLM。
• 保存和加载模型权重(model weights)，以便继续训练 LLM。
• 加载 OpenAI 提供的预训练权重(pretrained weights)。

在前几章中，我们已经完成了：

• 数据采样(data sampling)。
• 注意力机制(attention mechanism) 的实现。
• LLM 架构的编写。

本章的核心重点是：

• 实现训练函数(training function)。
• 预训练 LLM(pretrain the LLM)，如 图 5.1 所示。

图 5.1 说明 LLM 开发的三个主要阶段：1. 编写 LLM 代码（已完成）；2. 在通用文本数据集上预训练 LLM（本章重点）；3. 在有标注数据集上微调 LLM(finetuning)。本章专注于 预训练阶段，其中包括：1. 实现训练代码(training code)；2. 评估模型性能(evaluating performance)；3. 保存和加载模型权重(saving & loading model weights)。

正如图 5.1 所示，我们还将学习基本的模型评估技术，以衡量 LLM 生成文本的质量，这对于优化训练过程中的 LLM 至关重要。此外，我们将讨论如何加载预训练权重(pretrained weights)，为 LLM 提供一个良好的微调起点，以便在后续章节中进行微调(finetuning)。

权重参数(Weight Parameters)

在 LLM 及其他深度学习模型 中：

• 权重(weights) 指的是 训练过程中调整的可训练参数(trainable parameters)。
• 这些权重也被称为“权重参数(weight parameters)”或简单地称为“参数(parameters)”。

在 PyTorch 等深度学习框架 中：

• 权重存储在线性层(linear layers) 中，例如：在 第 3 章 实现的多头注意力模块(multi-head attention module) 和 第 4 章 实现的 GPTModel。
• 在 初始化一个层 后（如 new_layer=torch.nn.Linear(...)），
  可以通过 .weight 属性 访问其权重，即 new_layer.weight。

此外：

• PyTorch 允许直接访问模型的所有可训练参数(包括权重和偏置)，
  可通过 model.parameters() 方法获取。
• 在后续实现模型训练(model training) 时，我们将使用 model.parameters() 方法。

5.1 评估文本生成模型

本章的第一步是设置 LLM 进行文本生成，该部分基于前一章的代码进行扩展。在本节中，我们将讨论评估生成文本质量的基本方法。本节及本章其余部分的内容概览如图 5.2 所示。

图 5.2 展示了本章的核心内容，我们将先回顾上一章的文本生成过程，然后实现用于预训练阶段的基本模型评估技术，以衡量 LLM 生成文本的质量。

如图 5.2 所示，接下来的小节将回顾上一章末尾设置的文本生成过程，然后在后续小节深入探讨文本评估以及训练和验证损失的计算。

5.1.1 使用 GPT 生成文本

本节将设置 LLM 并简要回顾第 4 章实现的文本生成过程。我们首先初始化将在本章中评估和训练的 GPT 模型，该模型使用 GPTModel 类 和 GPT_CONFIG_124M 配置字典（来自第 4 章）：

import torch
'''记得修改路径'''
from chapter04 import GPTModel

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

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

在 GPT_CONFIG_124M 配置字典中，相比于上一章，我们唯一的调整是将上下文长度(context_length) 降至 256 个 token。这一修改降低了模型训练的计算需求，使得训练可以在标准笔记本电脑上运行。

最初，具有 1.24 亿参数 的 GPT-2 模型被配置为最多处理 1,024 个 token。在本章的训练过程结束后，我们将更新上下文长度(context size) 设置，并加载预训练权重(pretrained weights)，使模型能够处理 1,024-token 的上下文长度。

在 GPTModel 实例的基础上，我们使用上一章介绍的 generate_text_simple 函数，并引入两个实用函数：

• text_to_token_ids：用于将文本转换为 token ID。
• token_ids_to_text：用于将 token ID 还原为文本。

这两个函数能够在文本和 token 表示之间进行转换，我们将在本章反复使用这一技术。为了更直观地理解这一过程，图 5.3 进行了说明，随后我们将进入具体代码实现。

图 5.3 说明：文本生成(text generation) 过程包括：1. 将文本编码为 token ID，供 LLM 处理；2. LLM 计算 logit 向量(logit vectors)；3. 将 logit 向量转换回 token ID，并解码(detokenize) 为文本表示。

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

1. 分词器(tokenizer) 将输入文本转换为 token ID 序列（第 2 章已介绍）。
2. 模型接收 token ID，并生成对应的 logits（即表示词汇表中每个 token 概率分布的向量，第 4 章已介绍）。
3. 将 logits 转换回 token ID，并由分词器解码(detokenize) 为可读文本，完成从文本输入到文本输出的循环。

代码实现：文本生成过程：

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)  # 添加 batch 维度
    return encoded_tensor

def token_ids_to_text(token_ids, tokenizer):
    flat = token_ids.squeeze(0)  # 移除 batch 维度
    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

根据输出结果可以看出，由于模型尚未经过训练，它仍无法生成连贯的文本。为了定义何谓“连贯”或“高质量”文本，我们需要实现一种数值评估方法，用于衡量生成内容的质量。这样可以在整个训练过程中监控和优化模型的性能。

接下来的部分将介绍如何计算损失(loss) 作为评估生成输出的指标。
该损失函数将用于衡量训练进度和成功率。此外，在后续LLM 微调(finetuning) 章节中，我们还将探讨更多评估模型质量的方法。

5.1.2 计算文本生成损失

本节探讨如何通过计算文本生成损失(text generation loss) 来数值化评估模型在训练过程中生成的文本质量。我们将通过逐步讲解的方式，结合实际示例，使这些概念清晰且可应用。首先，我们回顾 第 2 章如何加载数据，以及如何使用第 4 章的 generate_text_simple 函数生成文本。

图 5.4 展示了 从输入文本到 LLM 生成文本的整体流程，共 五个步骤：1. 对于每个输入 token（左侧的 3 个 token），计算一个概率向量，该向量对应于词汇表中的所有 token；2. 在每个向量中，找到概率最高的索引位置，该索引对应于模型预测的下一个最可能 token ID；3. 选取这些概率最高的 token ID，并将其映射回文本；4. 生成的文本作为模型输出，代表 LLM 生成的文本内容；5. 该过程可持续进行，形成完整的文本序列。

文本生成过程 在 图 5.4 中展示了 generate_text_simple 函数（第 4 章介绍）的内部工作机制。在计算衡量生成文本质量的损失(loss) 之前，我们需要执行相同的初始步骤。

图 5.4 采用了仅包含 7 个 token 的小型词汇表（实际上就是每个token的维度，也就是用多少个维度来表示token的特征），以便图像能够适应单页展示。然而，我们的 GPTModel 采用的是 更大规模的词汇表（共 50,257 个词），因此，以下代码中的 token ID 范围为 0 到 50,256，而非 0 到 6。

此外，图 5.4 仅展示了单个示例（“every effort moves”），而下面的代码示例使用两个输入文本（“every effort moves” 和 “I really like”）作为 GPT 模型的输入：

输入示例（已映射到 token ID，对应于图 5.4 第 1 步）：

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

目标值 targets（即期望的模型输出 token ID）：

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

请注意，targets 实际上是 inputs 向右偏移 1 个位置的结果。这种 位移策略（shifting strategy） 是 关键步骤，用于训练模型预测序列中的下一个 token（在 第 2 章 的数据加载器实现中已介绍）。

当我们将 inputs 输入到模型中以计算两个输入示例的 logit 向量（每个示例包含三个 token）时，然后应用 softmax 函数 将这些 logit 值转换为 概率分数(probability scores)，这对应于 图 5.4 的第 2 步：

with torch.no_grad():  # A
    logits = model(inputs)
probas = torch.softmax(logits, dim=-1)  # 计算词汇表中每个 token 的概率
print(probas.shape)

概率张量 probas 维度如下：

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

• 第一维 (2)：对应于 两个输入示例（即 batch size）。
• 第二维 (3)：对应于 每个输入示例中的 token 数量。
• 第三维 (50257)：对应于 词汇表大小，即 token 的嵌入维度（embedding dimensionality）。

在通过 softmax 函数将 logits 转换为概率分数后，第 4 章 的 generate_text_simple 函数 会将这些概率分数转换回文本，对应于 图 5.4 的步骤 3-5。

我们可以通过 argmax 函数 对 概率分数 进行处理，以获取相应的 token ID（即 步骤 3 和 4）：

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

鉴于我们有2 个输入批次，每个批次包含3 个 token，对概率分数应用 argmax 函数（对应于 图 5.4 的第 3 步）后，得到2 组输出，每组包含 3 个预测 token ID：

Token IDs:
 tensor([[[16657],  # 第一批次
          [  339],
          [42826]],

         [[49906],  # 第二批次
          [29669],
          [41751]]])

最终， 第 5 步 将 token ID 转换回文本：

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)}")

解码后，我们发现这些输出 token 与目标 token 差别很大：

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

模型生成的文本是随机的，与目标文本不同，原因是模型尚未经过训练。现在，我们进入数值评估模型生成文本性能的部分，这种评估方法称为“损失”(loss)，如 图 5.4 所示。损失不仅用于衡量生成文本的质量，还是训练函数的重要组成部分，后续我们将使用它来更新模型权重，以改进生成文本。

图 5.5：在本节的剩余部分，我们将实现文本评估函数。在下一节，我们会将该评估函数应用到整个训练数据集。

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

模型训练的目标是提高 softmax 在正确目标 token ID 对应索引位置的概率，如 图 5.6 所示。这一 softmax 概率也用于本节剩余部分要实现的评估指标，用于数值评估模型生成的输出：正确位置的概率越高，模型生成的文本质量越好。

图 5.6：在训练前，模型生成随机的下一个 token 概率向量。模型训练的目标是确保与目标 token ID 对应的概率值被最大化。

请记住，图 5.6 显示的是仅包含 7 个 token 的紧凑词汇表的 softmax 概率，以便能够将所有内容适当地展示在一个图中。这意味着初始的随机值大约为 $1/7\approx 0.14$。

然而，在我们用于 GPT-2 模型 的词汇表中，共包含 50,257 个 token，因此大多数初始概率值会围绕 $1/50,257\approx 0.00002$ 波动。

对于两个输入文本，我们可以使用以下代码打印对应于目标 token 的初始 softmax 概率分数：

text_idx = 0
target_probas_1 = probas[text_idx, [0, 1, 2], targets[text_idx]]
print("Text 1:", target_probas_1)

text_idx = 1
target_probas_2 = probas[text_idx, [0, 1, 2], targets[text_idx]]
print("Text 2:", target_probas_2)

每个 batch 的 3 个目标 token ID 的概率 结果如下：

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

训练 LLM 的目标是最大化这些值，使其尽可能接近概率 1。这样，我们可以确保 LLM 在生成文本时始终选择正确的目标 token（即句子中的下一个单词）作为下一个生成的 token。

反向传播（Backpropagation）：

我们如何最大化 softmax 概率值，使其对应于目标 token？总体思路是更新模型权重，使模型对我们希望生成的目标 token ID 输出更高的值。权重更新是通过反向传播（backpropagation） 完成的，这是一种训练深度神经网络的标准技术（关于反向传播和模型训练的更多详细内容，请参考附录 A 中的A.3 到 A.7）。

反向传播需要一个损失函数（loss function），用于计算模型的预测输出（即目标 token ID 对应的概率）与真实目标输出之间的差异。这个损失函数衡量了模型的预测值与目标值之间的误差。

在本节的剩余部分，我们将计算两个示例 batch（target_probas_1 和 target_probas_2）的概率分数的损失值。 主要步骤如图 5.7 所示。

图 5.7 计算损失的过程涉及多个步骤：1. 步骤 1 至 3 计算与目标张量（target tensors）对应的 token 概率；2. 步骤 4 至 6 通过对概率值取对数并取平均值，完成损失计算。

由于我们已经完成了图 5.7 所列的步骤 1-3，即获得了 target_probas_1 和 target_probas_2，接下来我们进行步骤 4，对概率得分取对数：

log_probas = torch.log(torch.cat((target_probas_1, target_probas_2)))
print(log_probas)

这将产生以下数值：

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

在数学优化中，直接处理对数概率分数（log probability scores）比直接处理概率值更为便捷。这个主题超出了本书的范围，但我在附录 B 的参考资料中提供了一次讲座的链接，详细介绍了该主题。

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

avg_log_probas = torch.mean(log_probas)
print(avg_log_probas)

计算得到的平均对数概率如下：

tensor(-10.7940)

我们的目标是通过更新模型权重，使平均对数概率尽可能接近 0，这一过程将在第 5.2 节 进行实现。

然而，在深度学习中，通常的做法不是将平均对数概率提升到 0，而是将负的平均对数概率降到 0。负的平均对数概率计算如下（对应图 5.7 的步骤 6）：

neg_avg_log_probas = avg_log_probas * -1
print(neg_avg_log_probas)

输出：

tensor(10.7940)

在深度学习中，这个值**（-10.7940 取负变成 10.7940）被称为交叉熵损失（cross entropy loss）**。

PyTorch 提供了一个内置的 cross_entropy 函数，可以自动完成图 5.7 的所有 6 个步骤，大大简化了计算过程。

交叉熵损失（Cross entropy loss）：

交叉熵损失是机器学习和深度学习中常见的损失函数，用于衡量两个概率分布之间的差异——通常是真实标签的分布（在本例中是数据集中的真实 token）和模型预测的分布（例如LLM 生成的 token 概率）。

在机器学习，特别是 PyTorch 这样的框架中，cross_entropy 函数用于计算离散结果的交叉熵损失，这个计算过程与负的平均对数概率（negative average log probability）非常类似。因此，在实践中，交叉熵损失与负的平均对数概率这两个术语经常被互换使用。

在应用 cross_entropy 函数之前，我们先查看 logits 和 targets 张量的形状：

# Logits have shape (batch_size, num_tokens, vocab_size)
print("Logits shape:", logits.shape)

# Targets have shape (batch_size, num_tokens)
print("Targets shape:", targets.shape)

输出如下：

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

可以看到，logits 张量有 3 维：

• 第一维：批次大小（batch size）
• 第二维：每个样本的 token 数量
• 第三维：词汇表大小（vocabulary size）

而 targets 张量有 2 维：

• 第一维：批次大小
• 第二维：每个样本的 token 数量

在 PyTorch 的 cross_entropy_loss 函数中，我们需要将这些张量在批次维度上展平（flatten）：

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_flat 中的值是 LLM 需要生成的 token ID，而 logits_flat 包含未经过 softmax 处理的原始模型输出。

之前，我们手动进行了以下步骤：

1. 应用 softmax 计算概率分布
2. 选取目标 token ID 的概率
3. 计算负的平均对数概率

但是，PyTorch 的 cross_entropy 函数可以自动完成上述所有步骤，计算方式如下：

loss = torch.nn.functional.cross_entropy(logits_flat, targets_flat)
print(loss)

最终计算出的损失值与我们之前手动计算得到的结果完全相同（见图 5.7）：

tensor(10.7940)

困惑度（Perplexity）

困惑度是一种常用于语言建模（language modeling）的评估指标，通常与交叉熵损失（cross entropy loss）一起使用。它提供了一种更直观的方式，用于衡量模型在预测下一个 token 时的不确定性。

困惑度衡量模型预测的概率分布与数据集中实际词分布的匹配程度。与损失值（loss） 类似，较低的困惑度意味着模型的预测更接近真实分布。

困惑度的计算方式如下：

$$\text{perplexity}=\exp (\text{loss})$$

在 PyTorch 中，可以用 torch.exp(loss) 计算困惑度。对于我们之前计算的损失值，其结果如下：

perplexity = torch.exp(loss)
print(perplexity)

输出：

tensor(48725.8203)

困惑度比原始损失值更具可解释性，因为它表示模型在每个时间步中对多少个单词或 token 产生不确定性。在本例中，困惑度约为 48,725，这意味着模型在 48,725 个可能的单词或 token 之间不确定，不知道应该选择哪个作为下一个 token。

在本节中，我们针对两个小文本输入计算了损失（loss），用于演示如何评估模型的输出质量。在下一节，我们将把损失计算应用到整个训练集和验证集。

5.1.3 计算训练集和验证集的损失

在本节中，我们首先准备训练集和验证集，这些数据将在本章后续部分用于训练 LLM。然后，我们计算训练集和验证集的交叉熵损失（cross entropy loss），如图 5.8 所示。这是模型训练过程中的一个重要组成部分。

图 5.8 在上一节计算了交叉熵损失后，我们现在将其应用于整个文本数据集，该数据集将用于模型训练。

为了计算训练集和验证集的损失（如图 5.8 所示），我们使用了一个非常小的文本数据集——伊迪丝·华顿（Edith Wharton）撰写的短篇小说《裁决》（The Verdict）。我们在第 2 章已经使用过该文本。选择公共领域的文本可以避免与使用权相关的问题。此外，我们使用这样的小型数据集的原因是，它可以让代码示例在标准笔记本电脑上运行，即使没有高端 GPU，也可以在几分钟内完成训练，这对于教育目的来说特别有利。

对更大规模数据集感兴趣的读者也可以使用本书的附加代码，准备一个由超过 60,000 本 Project Gutenberg 公共领域书籍组成的数据集，并在此基础上训练一个 LLM（详见附录 D）。

预训练 LLM 的成本

为了让我们的项目规模更具现实感，我们可以参考Llama 2（一个相对流行的开放可用的 LLM）的训练成本。该模型拥有 70 亿个参数，训练过程中使用了184,320 GPU 小时，并处理了2 万亿个 token。

在撰写本文时，在 AWS 上运行 8xA100 GPU 云服务器的费用大约为 $30/小时。据此进行粗略估算，该 LLM 的总训练成本大约为 $690,000（计算方式：$184,320\text{小时}÷8\times 30$）。

以下代码用于加载我们在第 2 章使用过的短篇小说 The Verdict：

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"
    # 从指定URL下载文件并保存为"the-verdict.txt"
    urllib.request.urlretrieve(url, file_path)

file_path = "the-verdict.txt"
with open(file_path, "r", encoding="utf-8") as file:
    text_data = file.read()

加载数据集后，我们可以检查数据集中字符数和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 章的数据加载器（data loaders）来准备批量数据，以便用于 LLM 训练。该过程如图 5.9 所示。

图 5.9 说明 在准备数据加载器时，我们首先将输入文本拆分为训练集和验证集。然后，我们对文本进行标记化（图中仅展示了训练集部分，以简化表示）。接着，我们将标记化后的文本按照用户指定的长度（此处为 6）进行切块。最后，我们对这些切块进行随机排列，并按照批量大小（此处为 2）组织数据，以便用于模型训练。

为了便于可视化，图 5.9 采用了 max_length=6，这是由于空间限制。然而，在我们实际实现的数据加载器中，我们将 max_length 设置为 256-token，即 LLM 所支持的上下文长度，以便 LLM 在训练过程中能够看到更长的文本。

不同长度的训练
我们在训练过程中使用大小相似的文本块来简化处理并提高效率。然而，在实际应用中，使用可变长度输入训练 LLM 可能更有利，因为这有助于 LLM 更好地泛化，适用于不同类型的输入数据。

为了实现图 5.9 中的数据拆分和加载，我们首先定义 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 数据子集，我们可以调用 chapter 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
)

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
)

在上面的代码中，我们选择了较小的 batch_size，以降低计算资源的需求，因为我们正在使用非常小的数据集。然而，在实际训练 LLM 时，batch_size=1024 或更大是较为常见的。

作为可选的检查步骤，我们可以遍历数据加载器，确保它们被正确创建：

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])

从上述代码输出可以看出，训练集中有 9 个 batch，每个 batch 包含 2 个样本，每个样本 256 个 token。由于我们仅划分 10% 的数据用于验证集，因此只有 1 个验证 batch，其中包含 2 个输入样本。

正如第 2 章所讨论的，输入数据 (x) 和目标数据 (y) 具有相同的形状（batch_size × token数），因为目标数据只是输入数据右移一位。

接下来，我们实现一个计算单个 batch 交叉熵损失的实用函数 calc_loss_batch：

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

现在，我们可以使用 calc_loss_batch 计算单个 batch 的损失，然后实现 calc_loss_loader 函数，用于计算整个数据集的损失：

代码清单 5.2 计算训练和验证损失的函数

def calc_loss_loader(data_loader,model,device,num_batches=None):
    total_loss=0.
    if num_batches is None:
        num_batches=len(data_loader) #A
    else:
        num_batches=min(num_batches,len(data_loader)) #B
    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() #C
        else:
            break
    return total_loss/num_batches #D

默认情况下，calc_loss_loader 函数会遍历所有 batch 并将损失累加到 total_loss 变量中，最后计算并返回平均损失。另外，我们可以通过 num_batches 指定较小的 batch 数量，以便在模型训练期间加快评估过程。

计算训练集和验证集的损失
接下来，我们调用 calc_loss_loader 函数来计算训练集和验证集的损失：

device=torch.device("cuda" if torch.cuda.is_available() else "cpu") #A
model.to(device)
train_loss=calc_loss_loader(train_loader,model,device) #B
val_loss=calc_loss_loader(val_loader,model,device)
print("Training loss:",train_loss)
print("Validation loss:",val_loss)

计算结果

Training loss: 10.987583266364204
Validation loss: 10.98110580444336

损失值较高是因为模型尚未训练。作为对比，如果模型能够正确预测训练集和验证集中的下一个 token，则损失会接近 0。

现在我们已经有了一种方法来衡量生成文本的质量，在接下来的章节中，我们将训练 LLM 以降低损失值，使其更擅长文本生成，如图 5.10 所示。

图 5.10 我们已经回顾了文本生成过程，并实现了基本的模型评估技术来计算训练集和验证集的损失。接下来，我们将进入训练函数的实现，并对 LLM 进行预训练。

5.2 训练 LLM

在本节中，我们最终实现 LLM（我们的 GPTModel）的预训练代码。为此，我们专注于一个简洁易读的训练循环，如图 5.11 所示。然而，感兴趣的读者可以在附录 D（为训练循环添加额外技巧）中了解更高级的训练技术，包括学习率预热（learning rate warmup）、余弦退火（cosine annealing）和梯度裁剪（gradient clipping）。

图 5.11 一个用于 PyTorch 深度神经网络训练的典型训练循环，包括多个步骤，针对训练集中的小批量（batch）数据进行多轮（epoch）迭代。在每个循环中，我们计算每个训练集批次的损失（loss），以此确定损失梯度（loss gradients），然后使用这些梯度更新模型权重（model weights），以最小化训练集损失。

图 5.11 的流程图展示了一个典型的 PyTorch 神经网络训练工作流程，我们将在训练 LLM 时采用相同的方法。该流程概述了八个步骤，包括：

1. 遍历所有训练轮数（epoch）
2. 处理小批量数据（batch）
3. 重置梯度（reset gradients）
4. 计算梯度（compute gradients）
5. 更新权重（update weights）
6. 监控训练（monitor training）
7. 打印损失（print losses）
8. 生成文本样本（generate text samples）

如果你对 PyTorch 训练深度神经网络不太熟悉，并且这些步骤对你来说还不够清楚，建议阅读附录 A（PyTorch 入门）的A.5 到 A.8 小节。

在代码中，我们可以通过以下 train_model_simple 函数实现该训练流程：

代码清单 5.3：LLM 预训练的主函数

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 = [], [], []  # A
    tokens_seen, global_step = 0, -1

    # Main training loop
    for epoch in range(num_epochs):  # B
        model.train()  # Set model to training mode
        
        for input_batch, target_batch in train_loader:
            optimizer.zero_grad()  # C # Reset loss gradients from previous batch iteration
            loss = calc_loss_batch(input_batch, target_batch, model, device)
            loss.backward()  # D # Calculate loss gradients
            optimizer.step()  # E # Update model weights using loss gradients
            '''
            input_batch.numel() 返回批次中元素总数。
            例如，若 input_batch 形状为 (32, 128)，则 numel() = 32×128 = 4096。
            将这个数累加到 tokens_seen 中。
            '''
            tokens_seen += input_batch.numel()
            global_step += 1

            # Optional evaluation step
            if global_step % eval_freq == 0:  # F
                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(  # G
            model, tokenizer, device, start_context
        )

    return train_losses, val_losses, track_tokens_seen

需要注意的是，我们刚刚创建的 train_model_simple 函数依赖于两个尚未定义的函数：

• evaluate_model
• generate_and_print_sample

其中，evaluate_model 对应于 图 5.11 中的步骤 7。该函数会在每次模型更新后打印训练集和验证集的损失值，以便我们评估训练是否提高了模型的性能。

更具体地说，evaluate_model 函数在计算训练集和验证集上的损失时，会确保模型处于评估模式，同时禁用梯度跟踪和 dropout：

def evaluate_model(model, train_loader, val_loader, device, eval_iter):
    model.eval()  # A
    with torch.no_grad():  # B
        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

与 evaluate_model 类似，generate_and_print_sample 也是一个便捷函数，用于跟踪模型在训练过程中是否得到改进。具体而言，generate_and_print_sample 函数接收一个文本片段 (start_context) 作为输入，将其转换为token ID，然后输入到 LLM 中，通过之前使用过的 generate_text_simple 函数生成文本样本：

def generate_and_print_sample(model, tokenizer, device, start_context):
    '''将模型设置为评估模式，关闭 dropout、batch normalization 的训练行为。'''
    model.eval()
    '''通过模型的 position embedding 层获取预设的上下文长度。'''
    context_size = model.pos_emb.weight.shape[0]
    '''将 start_context 字符串转换为 token IDs 序列，随后移动到指定设备。'''
    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()

虽然 evaluate_model 提供了训练进度的数值估计，但 generate_and_print_sample 提供了具体的文本示例，便于在训练过程中直观判断模型的生成能力。

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

让我们通过训练一个 GPTModel 实例 10 轮（epochs） 来实际演示这一过程，使用 AdamW 优化器，并调用之前定义的 train_model_simple 函数进行训练：

# Note:
# Uncomment the following code to calculate the execution time
# import time
# start_time = time.time()

torch.manual_seed(123)
model = GPTModel(GPT_CONFIG_124M)
model.to(device)
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
)

# Note:
# Uncomment the following code to show the execution time
# end_time = time.time()
# execution_time_minutes = (end_time - start_time) / 60
# print(f"Training completed in {execution_time_minutes:.2f} minutes.")

执行 train_model_simple 函数后，训练过程开始，在 MacBook Air 或类似的笔记本电脑上大约 5 分钟 内完成。执行过程中打印的输出如下：

Ep 1 (Step 000000): Train loss 9.817, Val loss 9.928
Ep 1 (Step 000005): Train loss 8.067, Val loss 8.334
Every effort moves you,,,,,,,,,,,,.                                     
[...] #A
Ep 10 (Step 000085): Train loss 0.625, Val loss 6.395
Every effort moves you?"  "Yes--quite insensible to the irony. She wanted him vindicated--and by me!"  He laughed again, and threw back his head to look up at the sketch of the donkey. "There were days when I

从训练过程中打印的结果可以看出，训练损失（Training Loss）显著下降，从初始值 9.817 收敛至 0.625。模型的语言能力得到了显著提升：

• 在训练开始时，模型只能在初始文本 “Every effort moves you” 之后随机添加逗号 (“Every effort moves you,”) 或者不断重复单词 “and”。
• 经过训练后，模型能够生成 符合语法的文本，如 “Every effort moves you?” “Yes–quite insensible to the irony…”。

与训练集损失类似，验证集损失（Validation Loss） 也从最初的 9.928 逐渐下降。然而，与训练集损失不同，验证集损失并未像训练损失一样显著下降，在 第 10 轮（epoch）后仍保持在 6.395。

在进一步讨论验证集损失之前，我们先绘制一个损失曲线图，以便直观比较训练损失和验证损失的变化趋势：

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)

运行上述代码后，生成的训练损失与验证损失曲线图 如 图 5.12 所示。

图 5.12 显示，在训练初期，训练集和验证集的损失均迅速下降，这表明模型正在学习。然而，在 第二轮（epoch 2） 之后，训练集损失仍然持续下降，而验证集损失却趋于停滞。这表明模型仍在学习，但 从第 2 轮开始，它已经过拟合（overfitting）到训练集。

如 图 5.12 所示，训练损失和验证损失在 第一轮（epoch 1） 内均有所下降。然而，在 第二轮（epoch 2） 之后，两者开始出现分歧。这种 损失的分歧 以及 验证损失明显高于训练损失 的现象表明，模型正在过拟合（overfitting）训练数据。

我们可以通过在 The Verdict 文本文件中搜索模型生成的文本片段（例如 “quite insensible to the irony”）来验证 模型是否在逐字记忆（memorizing）训练数据。

这种记忆效应是可以预期的，因为我们使用了 非常小的训练数据集 并 对模型训练了多个轮次。通常，在 大规模数据集上，模型只会训练一个 epoch，而不是多次遍历相同的数据集。如前所述，感兴趣的读者可以尝试使用 Project Gutenberg 的 60,000 本公共领域书籍 训练该模型，在这种情况下就不会出现过拟合的情况；详见 附录 B。

在接下来的章节中（见 图 5.13），我们将探讨 LLMs 采用的采样方法，这些方法有助于缓解 记忆效应，从而生成更加新颖的文本。

如 图 5.13 所示，在实现训练函数后，我们的模型已经能够生成 连贯的文本。然而，模型仍然经常逐字记忆训练集中的段落。接下来的章节将介绍 生成更具多样性文本的策略。

如 图 5.13 所示，下一节将介绍 LLM 的文本生成策略，以 减少训练数据的记忆 并 提高 LLM 生成文本的原创性。在此之后，我们将探讨 权重的保存与加载，以及 从 OpenAI 的 GPT 模型加载预训练权重。

5.3 解码策略以控制随机性

在本节中，我们将介绍 文本生成策略（也称为 解码策略），以 生成更具原创性的文本。首先，我们 回顾 之前章节中的 generate_text_simple 函数，该函数已在本章早些时候的 generate_and_print_sample 中使用。然后，我们将介绍 两种改进该函数的技术：温度缩放（temperature scaling） 和 Top-k 采样（top-k sampling）。

我们首先 将模型从 GPU 迁移回 CPU，因为对于相对较小的模型来说，在推理过程中 不需要 GPU。此外，在 训练完成后，我们将模型设置为 评估模式（evaluation mode），以 关闭 dropout 等随机组件：

model.to("cpu")
model.eval()

接下来，我们将 GPTModel 实例（model）传入 generate_text_simple 函数，该函数 使用 LLM 逐个生成 token：

tokenizer = tiktoken.get_encoding("gpt2")

token_ids = generate_text_simple(
    model=model,
    idx=text_to_token_ids("Every effort moves you", tokenizer),
    max_new_tokens=25,
    context_size=GPT_CONFIG_124M["context_length"]
)

print("Output text:\n", token_ids_to_text(token_ids, tokenizer))

生成的文本如下：

Output text:
 Every effort moves you?"

"Yes--quite insensible to the irony. She wanted him vindicated--and by me!"

如 5.1.2 节 所述，每次生成 下一个 token 时，LLM 会选择 所有词汇表 token 中概率最高的 token。

这意味着 LLM 在相同输入上下文（“Every effort moves you”）下，每次运行 generate_text_simple 函数时，都会生成完全相同的文本。

接下来的小节将介绍 两种方法 来 控制生成文本的随机性和多样性：

1. 温度缩放（temperature scaling）
2. Top-k 采样（top-k sampling）

5.3.1 温度缩放（Temperature Scaling）

本节介绍了温度缩放技术，它在下一个 token 的生成任务中引入了概率选择过程。

在之前的 generate_text_simple 函数中，我们总是使用 torch.argmax 选择概率最高的 token 作为下一个 token，这种方式也被称为贪婪解码（greedy decoding）。为了生成更具多样性的文本，我们可以将 argmax 替换为从概率分布中采样的函数（也就是从 LLM 在每一步生成时为词表中每个词生成的概率中采样）。

为了用一个具体的例子来说明概率采样的过程，我们暂且使用一个非常小的词表进行说明：

vocab = {
    "closer": 0,
    "every": 1,
    "effort": 2,
    "forward": 3,
    "inches": 4,
    "moves": 5,
    "pizza": 6,
    "toward": 7,
    "you": 8,
}
inverse_vocab = {v: k for k, v in vocab.items()}

接下来，假设 LLM 给定的起始上下文是 "every effort moves you"，并且生成了以下的下一个 token 的 logits：

next_token_logits = torch.tensor(
    [4.51, 0.89, -1.90, 6.75, 1.63, -1.62, -1.89, 6.28, 1.79]
)

如上一章所述，在 generate_text_simple 函数内部，我们使用 softmax 函数将 logits 转换为概率分布，并使用 argmax 函数选择概率最高的 token 的 ID，然后使用反向词表将其映射回文本：

probas = torch.softmax(next_token_logits, dim=0)
next_token_id = torch.argmax(probas).item()
print(inverse_vocab[next_token_id])

由于 logits 值中最大的值在第四个位置（索引位置为3，因为 Python 使用从0开始的索引），因此生成的单词是 "forward"。

为了实现一个概率采样过程，我们现在可以将 argmax 替换为 PyTorch 中的 multinomial 函数：

torch.manual_seed(123)
next_token_id = torch.multinomial(probas, num_samples=1).item()
print(inverse_vocab[next_token_id])

打印的输出是 "forward"，和之前一样。发生了什么呢？multinomial 函数根据概率分数按比例采样下一个 token。换句话说，"forward" 仍然是最可能的 token，因此它会在大多数时候被 multinomial 选中，但并不是每一次都会选中。为了说明这一点，我们来实现一个函数，将采样过程重复 1000 次：

def print_sampled_tokens(probas):
    torch.manual_seed(123) # Manual seed for reproducibility
    sample = [torch.multinomial(probas, num_samples=1).item() for i in range(1_000)]
    sampled_ids = torch.bincount(torch.tensor(sample), minlength=len(probas))
    for i, freq in enumerate(sampled_ids):
        print(f"{freq} x {inverse_vocab[i]}")

print_sampled_tokens(probas)

采样输出如下：

71 x closer
2 x every
0 x effort
544 x forward
2 x inches
1 x moves
0 x pizza
376 x toward
4 x you

正如我们从输出中所看到的，单词 "forward" 在大多数时候被采样（1000 次中有 582 次），但其他的 token，比如 "closer"、"inches" 和 "toward" 也会在一些时候被采样。这意味着，如果我们将 generate_and_print_sample 函数中的 argmax 替换为 multinomial 函数，LLM 有时就会生成如下的文本：

• "every effort moves you toward"
• "every effort moves you inches"
• "every effort moves you closer"

而不是始终生成 "every effort moves you forward"。

我们还可以通过一个叫作温度缩放（temperature scaling）的概念进一步控制分布和选择过程，温度缩放其实就是一个看起来高级的说法，表示将 logits 除以一个大于 0 的数字：

def softmax_with_temperature(logits, temperature):
    scaled_logits = logits / temperature
    return torch.softmax(scaled_logits, dim=0)

当温度大于 1 时，token 的概率分布将变得更均匀；当温度小于 1 时，分布将变得更加自信（更尖锐或更陡峭）。让我们通过绘图来展示原始概率与不同温度缩放后的概率之间的差异：

temperatures = [1, 0.1, 5]  # 原始温度、更低温度和更高温度
scaled_probas = [softmax_with_temperature(next_token_logits, T) for T in temperatures]
x = torch.arange(len(vocab))
bar_width = 0.15
fig, ax = plt.subplots(figsize=(5, 3))
for i, T in enumerate(temperatures):
    rects = ax.bar(x + i * bar_width, scaled_probas[i],
                   bar_width, label=f'Temperature = {T}')
ax.set_ylabel('Probability')
ax.set_xticks(x)
ax.set_xticklabels(vocab.keys(), rotation=90)
ax.legend()
plt.tight_layout()
plt.show()

print_sampled_tokens(scaled_probas[1])

0 x closer
0 x every
0 x effort
992 x forward
0 x inches
0 x moves
0 x pizza
8 x toward
0 x you

print_sampled_tokens(scaled_probas[2])

153 x closer
68 x every
55 x effort
223 x forward
102 x inches
50 x moves
43 x pizza
218 x toward
88 x you

绘制出的图如图 5.14 所示。

图 5.14 表示温度为 1 时是词表中每个 token 的未缩放概率分布。当温度降低至 0.1 时，分布变得更加陡峭，因此最可能的 token（此处为 “forward”）的概率分数将进一步提高。相反，温度升高至 5 会使概率分布更加均匀。

当温度为 1 时，会在将 logits 传递给 softmax 函数以计算概率分数之前将 logits 除以 1。换句话说，使用温度为 1 与不使用任何温度缩放是等价的。在这种情况下，tokens 会通过 PyTorch 的 multinomial 采样函数，以等于原始 softmax 概率分数的概率被选中。

例如，对于温度设置为 1 的情况，如图 5.14 所示，“forward” 对应的 token 会在大约 60% 的时间内被选中。

同时，如图 5.14 所示，应用非常小的温度（如 0.1）将导致分布更加陡峭，因此 multinomial 函数的行为会几乎在 100% 的时间内选择最可能的 token（此处是 “forward”），这趋近于 argmax 函数的行为。反之，温度为 5 会导致分布更加均匀，其他 token 更频繁地被选择。这可以为生成的文本添加更多的多样性，但也更可能导致无意义的文本。例如，使用温度为 5 时，大约有 4% 的概率生成诸如 “every effort moves you pizza” 这样的文本。

练习 5.1

使用 print_sampled_tokens 函数，打印出使用图 5.13 中所示温度缩放后的 softmax 概率的采样频率。在每种情况下，单词 “pizza” 被采样的频率是多少？你能否想到一种更快、更精确的方法来确定 “pizza” 被采样的频率？

import torch

vocab = { 
    "closer": 0,
    "every": 1, 
    "effort": 2, 
    "forward": 3,
    "inches": 4,
    "moves": 5, 
    "pizza": 6,
    "toward": 7,
    "you": 8,
} 
inverse_vocab = {v: k for k, v in vocab.items()}

next_token_logits = torch.tensor(
    [4.51, 0.89, -1.90, 6.75, 1.63, -1.62, -1.89, 6.28, 1.79]
)

def print_sampled_tokens(probas):
    torch.manual_seed(123)
    sample = [torch.multinomial(probas, num_samples=1).item() for i in range(1_000)]
    sampled_ids = torch.bincount(torch.tensor(sample))
    for i, freq in enumerate(sampled_ids):
        print(f"{freq} x {inverse_vocab[i]}")


def softmax_with_temperature(logits, temperature):
    scaled_logits = logits / temperature
    return torch.softmax(scaled_logits, dim=0)


temperatures = [1, 0.1, 5]  # Original, higher, and lower temperature
scaled_probas = [softmax_with_temperature(next_token_logits, T) for T in temperatures]

for i, probas in enumerate(scaled_probas):
    print("\n\nTemperature:", temperatures[i])
    print_sampled_tokens(probas)


temp5_idx = 2
pizza_idx = 6

print(scaled_probas[temp5_idx][pizza_idx])

Temperature: 1
71 x closer
2 x every
0 x effort
544 x forward
2 x inches
1 x moves
0 x pizza
376 x toward
4 x you


Temperature: 0.1
0 x closer
0 x every
0 x effort
992 x forward
0 x inches
0 x moves
0 x pizza
8 x toward


Temperature: 5
153 x closer
68 x every
55 x effort
223 x forward
102 x inches
50 x moves
43 x pizza
218 x toward
88 x you

tensor(0.0430)

5.3.2 Top-k 采样

在前一节中，我们实现了一种结合温度缩放的概率采样方法，以增加输出的多样性。我们看到，较高的温度值会导致更均匀分布的下一个 token 概率，从而通过减少模型反复选择最可能 token 的可能性，生成更具多样性的输出。这种方法可以探索生成过程中不太可能但可能更有趣或更有创造性的路径。然而，这种方法的一个缺点是，它有时会导致语法不正确或完全无意义的输出，例如 “every effort moves you pizza”。

在本节中，我们引入另一个概念，称为 top-k 采样。当它与概率采样和温度缩放结合使用时，可以改进文本生成的结果。

在 top-k 采样中，我们可以将可采样的 token 限制为前 k 个最有可能的 token，并通过屏蔽它们的概率分数，将所有其他 token 排除在选择过程之外，如图 5.15 所示。

图 5.15 使用 $k=3$ 的 top-k 采样时，我们关注与最高 logits 值相关联的 3 个 token，并将所有其他 token 的 logits 设置为负无穷（$-\infty$）以屏蔽它们，然后再应用 softmax 函数。这将产生一个概率分布，其中所有非 top-k token 的概率值都为 0。

图 5.15 中所概述的方法将所有未被选中的 logit 替换为负无穷值（$-\infty$），这样在计算 softmax 值时，非 top-k token 的概率得分为 0，剩下的概率之和为 1。（细心的读者可能还记得我们在第 3 章第 3.5.1 节“应用因果注意力掩码”中实现的因果注意力模块中也使用过这种掩码技巧。）

在代码中，我们可以如下实现图 5.15 所示的 top-k 操作，首先选择具有最大 logit 值的 token：

top_k = 3
top_logits, top_pos = torch.topk(next_token_logits, top_k)

print("Top logits:", top_logits)
print("Top positions:", top_pos)

前 3 个 token 的 logit 值及其对应的 token ID（按降序排列）如下：

Top logits: tensor([6.7500, 6.2800, 4.5100])
Top positions: tensor([3, 7, 0])

随后，我们使用 PyTorch 的 where 函数，将不在 top-3 范围内的 token 的 logit 值设为负无穷（$-\infty$）：

new_logits = torch.where(
    condition=next_token_logits < top_logits[-1],
    input=torch.tensor(float("-inf")), 
    other=next_token_logits
)

print(new_logits)

下一个 token 的 9 个词汇表项的最终 logits 如下所示：

tensor([4.5100,   -inf,   -inf, 6.7500,   -inf,   -inf,   -inf, 6.2800,   -inf])

最后，让我们对这些 logits 应用 softmax 函数，将它们转换为下一个 token 的概率：

topk_probas = torch.softmax(new_logits, dim=0)
print(topk_probas)

如我们所见，这种 top-3 方法的结果是 3 个非零的概率值：

tensor([0.0615, 0.0000, 0.0000, 0.5775, 0.0000, 0.0000, 0.0000, 0.3610, 0.0000])

现在，我们可以应用上一节中介绍的温度缩放（temperature scaling）和多项分布采样（multinomial sampling）来在这 3 个非零概率中选择下一个 token，从而生成文本的下一个 token。我们将在下一节通过修改文本生成函数来实现这一过程。

5.3.3 修改文本生成函数

前两小节介绍了两个用于增加LLM生成文本多样性的概念：温度采样（temperature sampling）和top-k采样（top-k sampling）。在本节中，我们将这两个概念结合起来，并添加到我们之前用于通过LLM生成文本的generate_simple函数中，创建一个新的generate函数：

代码清单5.4 一个具有更多多样性的修改版文本生成函数：

def generate(model, idx, max_new_tokens, context_size, temperature=0.0, top_k=None, eos_id=None):
    """
    输入：
    - model：语言模型，接受一个 token 序列（通常为整数 ID 序列）输出对应的 logits（未归一化概率）。
    - idx：初始输入 token 序列，形状一般为 (batch_size, current_seq_length)。
    - max_new_tokens：要生成的新 token 数量上限。
    - context_size：模型每次考虑的上下文长度（从 idx 中取最后 context_size 个 token）。
    - temperature：温度参数，用于控制采样的随机性；温度大于 0 表示采样（随机性越大），等于 0 表示贪心取最大概率（确定性）。
    - top_k：如果指定，则进行 top-k 采样，即只保留 logits 中最大的 k 个，其余置为 -∞。
    - eos_id：可选的结束标识符，如果生成的 token 等于此 ID，则提前停止生成。
    输出：
    - 返回最终包含初始序列和生成新 token 的序列（token IDs）
    """
    # For-loop is the same as before: Get logits, and only focus on last time step
    '''循环最多生成 max_new_tokens 个 token。'''
    for _ in range(max_new_tokens):
        '''从当前序列 idx 中只取最后 context_size 个 token 作为当前模型的输入。'''
        idx_cond = idx[:, -context_size:]
        with torch.no_grad():
            '''得到 logits（未归一化概率分布）'''
            logits = model(idx_cond)
        '''取出 logits 中最后一个时间步对应的 logits，因为我们只关心当前生成的下一个 token。'''
        logits = logits[:, -1, :]

        # New: Filter logits with top_k sampling
        if top_k is not None:
            # Keep only top_k values
            '''对每个样本的 logits 取前 top_k 的最大值。'''
            top_logits, _ = torch.topk(logits, top_k)
            '''min_val 即是前 k 个中最小的一个（即第 k 大的值）。'''
            min_val = top_logits[:, -1]
            '''将 logits 中小于 min_val 的值设为 -∞，这样在 softmax 时其概率变为 0。'''
            logits = torch.where(logits < min_val, torch.tensor(float("-inf")).to(logits.device), logits)

        # New: Apply temperature scaling
        '''判断是否使用温度采样，温度大于 0 表示采用随机采样策略。'''
        if temperature > 0.0:
            '''
            作用：将 logits 除以 temperature 进行缩放。
            原因：温度控制采样分布的平滑程度，温度越低，分布越尖锐；温度越高，分布越平滑。
            '''
            logits = logits / temperature

            # Apply softmax to get probabilities
            '''
            作用：对缩放后的 logits 应用 softmax 函数，得到各 token 的概率分布。
            形状：输出形状 (batch_size, vocab_size)。
            '''
            probs = torch.softmax(logits, dim=-1)  # (batch_size, context_len)

            # Sample from the distribution
            '''
            作用：从概率分布中随机采样一个 token，multinomial 根据概率进行采样。
            输出：形状 (batch_size, 1)；例如可能得到 [[123]] 表示选择词表中第 123 个 token。
            '''
            idx_next = torch.multinomial(probs, num_samples=1)  # (batch_size, 1)

        # Otherwise same as before: get idx of the vocab entry with the highest logits value
        else:
            '''作用：若温度为 0，则不进行随机采样，而是直接选择 logits 中最大的那个 token（贪心选法）。'''
            idx_next = torch.argmax(logits, dim=-1, keepdim=True)  # (batch_size, 1)

        if idx_next == eos_id:  # Stop generating early if end-of-sequence token is encountered and eos_id is specified
            """若生成的 token 等于 eos_id（结束符号），则提前停止生成。"""
            break

        # Same as before: append sampled index to the running sequence
        '''
        作用：将新生成的 token（idx_next）拼接到当前序列 idx 的末尾。
        原因：保持生成序列的更新，后续生成会基于完整序列进行。
        '''
        idx = torch.cat((idx, idx_next), dim=1)  # (batch_size, num_tokens+1)

    return idx

现在我们来看看这个新的generate函数的实际效果：

torch.manual_seed(123)

token_ids = generate(
    model=model,
    idx=text_to_token_ids("Every effort moves you", tokenizer),
    max_new_tokens=15,
    context_size=GPT_CONFIG_124M["context_length"],
    top_k=25,
    temperature=1.4
)

print("Output text:\n", token_ids_to_text(token_ids, tokenizer))

生成的文本如下：

Output text:
 Every effort moves you know began to happen a little wild--I was such a good fellow enough

正如我们所见，这段生成的文本与我们在第5.3节开头通过generate_simple函数生成的文本非常不同（“Every effort moves you know," was one of the axioms he laid…!"），那段文本是训练集中记忆下来的片段。

练习5.2
尝试使用不同的温度和top-k设置。根据你的观察，你能想到哪些应用场景更适合使用较低的温度和top-k设置？反过来，你能想到哪些应用更适合使用较高的温度和top-k设置？（建议在本章最后加载OpenAI的预训练权重后再次完成此练习。）

温度和top-k设置都必须根据具体的LLM进行调整（这是一种反复试验的过程，直到它生成理想的输出为止）。
不过，理想的输出也是与具体应用相关的。
较低的top-k和温度设置会产生更少随机性的输出，这在创作教育内容、技术写作或问答系统、数据分析、代码生成等场景中是理想的。
较高的top-k和温度设置会产生更多样化和随机的输出，这在头脑风暴任务、创意写作等场景中更受欢迎。

练习5.3
对于generate函数，如何设置不同的参数组合来实现确定性行为，也就是说禁用随机采样，使其像generate_simple函数一样总是生成相同的输出？

设置为top_k=None并且不使用温度缩放；
设置top_k=1。

import tiktoken
import torch
from previous_chapters import GPTModel


GPT_CONFIG_124M = {
    "vocab_size": 50257,  # Vocabulary size
    "context_length": 256,       # Shortened context length (orig: 1024)
    "emb_dim": 768,       # Embedding dimension
    "n_heads": 12,        # Number of attention heads
    "n_layers": 12,       # Number of layers
    "drop_rate": 0.1,     # Dropout rate
    "qkv_bias": False     # Query-key-value bias
}


torch.manual_seed(123)

tokenizer = tiktoken.get_encoding("gpt2")
model = GPTModel(GPT_CONFIG_124M)
model.load_state_dict(torch.load("model.pth", weights_only=True))
model.eval();

from gpt_generate import generate, text_to_token_ids, token_ids_to_text
from previous_chapters import generate_text_simple
# Deterministic function that used torch.argmax

start_context = "Every effort moves you"

token_ids = generate_text_simple(
    model=model,
    idx=text_to_token_ids(start_context, tokenizer),
    max_new_tokens=25,
    context_size=GPT_CONFIG_124M["context_length"]
)

print("Output text:\n", token_ids_to_text(token_ids, tokenizer))

# Deterministic behavior: No top_k, no temperature scaling

token_ids = generate(
    model=model,
    idx=text_to_token_ids("Every effort moves you", tokenizer),
    max_new_tokens=25,
    context_size=GPT_CONFIG_124M["context_length"],
    top_k=None,
    temperature=0.0
)

print("Output text:\n", token_ids_to_text(token_ids, tokenizer))

# Deterministic behavior: No top_k, no temperature scaling

token_ids = generate(
    model=model,
    idx=text_to_token_ids("Every effort moves you", tokenizer),
    max_new_tokens=25,
    context_size=GPT_CONFIG_124M["context_length"],
    top_k=None,
    temperature=0.0
)

print("Output text:\n", token_ids_to_text(token_ids, tokenizer))

到目前为止，我们已经介绍了如何预训练LLMs并使用它们生成文本。本章最后两个部分将讨论如何保存和加载训练好的LLM，以及如何从OpenAI加载预训练权重。

5.4 在 PyTorch 中加载和保存模型权重

在本章中，我们已经讨论了如何通过数值方式评估训练进度，并从零开始预训练一个LLM。尽管该LLM和数据集都相对较小，这个练习仍然表明预训练LLM是计算资源密集型的。因此，能够保存LLM是很重要的，这样我们就不必在每次开启新的会话时重新运行训练。

如图5.16的章节概览所示，我们将在本节中介绍如何保存和加载一个预训练模型。随后，在下一节中，我们将把一个能力更强的OpenAI预训练GPT模型加载到我们的GPTModel实例中。

图5.16 在训练和检查模型之后，保存模型通常是很有帮助的，这样我们可以在之后使用或继续训练它。这正是本节的主题，在本章最后一节加载OpenAI提供的预训练模型权重之前我们会先完成这个步骤。

幸运的是，保存一个 PyTorch 模型相对简单。推荐的方式是使用 torch.save 函数保存模型的所谓 state_dict，它是一个将每一层映射到其参数的字典，如下所示：

torch.save(model.state_dict(), "model.pth")

在上述代码中，“model.pth” 是保存 state_dict 的文件名。.pth 扩展名是 PyTorch 文件的惯例，尽管从技术上讲我们也可以使用任何文件扩展名。

然后，在通过 state_dict 保存了模型权重之后，我们可以将这些权重加载到一个新的 GPTModel 模型实例中，如下所示：

model = GPTModel(GPT_CONFIG_124M)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.load_state_dict(torch.load("model.pth", map_location=device, weights_only=True))
model.eval()

如第4章所讨论的，dropout 通过在训练期间随机“丢弃”某一层的神经元来帮助防止模型对训练数据的过拟合。然而，在推理阶段，我们不希望随机丢弃网络已经学到的信息。使用 model.eval() 将模型切换到评估模式，以用于推理，这会禁用模型中的 dropout 层。

如果我们打算之后继续对模型进行预训练，例如使用本章前面定义的 train_model_simple 函数，建议也保存优化器的状态。

自适应优化器（如 AdamW）为每个模型权重存储额外的参数。AdamW 使用历史数据为每个模型参数动态调整学习率。如果没有这些历史信息，优化器会重置，模型可能会以次优的方式学习，甚至无法正确收敛，这意味着它将失去生成连贯文本的能力。我们可以使用 torch.save 同时保存模型和优化器的 state_dict 内容，如下所示：

torch.save({
    "model_state_dict": model.state_dict(),
    "optimizer_state_dict": optimizer.state_dict(),
    }, 
    "model_and_optimizer.pth"
)

然后，我们可以通过先使用 torch.load 加载保存的数据，再使用 load_state_dict 方法来恢复模型和优化器状态，如下所示：

checkpoint = torch.load("model_and_optimizer.pth", weights_only=True)

model = GPTModel(GPT_CONFIG_124M)
model.load_state_dict(checkpoint["model_state_dict"])

optimizer = torch.optim.AdamW(model.parameters(), lr=0.0005, weight_decay=0.1)
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
model.train()

练习 5.4
在保存权重之后，在一个新的 Python 会话或 Jupyter notebook 文件中加载模型和优化器，并使用 train_model_simple 函数继续对其进行 1 个 epoch 的预训练。

import tiktoken
import torch
from previous_chapters import GPTModel


GPT_CONFIG_124M = {
    "vocab_size": 50257,   # Vocabulary size
    "context_length": 256, # Shortened context length (orig: 1024)
    "emb_dim": 768,        # Embedding dimension
    "n_heads": 12,         # Number of attention heads
    "n_layers": 12,        # Number of layers
    "drop_rate": 0.1,      # Dropout rate
    "qkv_bias": False      # Query-key-value bias
}

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

tokenizer = tiktoken.get_encoding("gpt2")

checkpoint = torch.load("model_and_optimizer.pth", weights_only=True)
model = GPTModel(GPT_CONFIG_124M)
model.load_state_dict(checkpoint["model_state_dict"])
model.to(device)

optimizer = torch.optim.AdamW(model.parameters(), lr=0.0004, weight_decay=0.1)
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
model.train();

import os
import urllib.request
from previous_chapters import create_dataloader_v1


file_path = "the-verdict.txt"
url = "https://raw.githubusercontent.com/rasbt/LLMs-from-scratch/main/ch02/01_main-chapter-code/the-verdict.txt"

if not os.path.exists(file_path):
    with urllib.request.urlopen(url) as response:
        text_data = response.read().decode('utf-8')
    with open(file_path, "w", encoding="utf-8") as file:
        file.write(text_data)
else:
    with open(file_path, "r", encoding="utf-8") as file:
        text_data = file.read()


# Train/validation ratio
train_ratio = 0.90
split_idx = int(train_ratio * len(text_data))
train_data = text_data[:split_idx]
val_data = text_data[split_idx:]


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
)

from gpt_train import train_model_simple

num_epochs = 1
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
)

5.5 从 OpenAI 加载预训练权重

此前，出于教学目的，我们使用一个由短篇小说组成的小型数据集训练了一个小型的 GPT-2 模型。这种方法使我们能够专注于基础知识，而无需投入大量时间和计算资源。

幸运的是，OpenAI 公开分享了其 GPT-2 模型的权重，从而省去了我们自己在大型语料库上重新训练模型所需的数万甚至数十万美元的投入。

在本节的其余部分中，我们将把这些权重加载到我们自己的 GPTModel 类中，并使用该模型进行文本生成。这里的“权重”是指存储在 PyTorch 的 Linear 和 Embedding 层的 .weight 属性中的可训练参数。例如，在训练模型时我们曾通过 model.parameters() 来访问它们。

在接下来的章节中，我们将重用这些预训练权重，对模型进行微调，使其能够执行文本分类任务，并像 ChatGPT 一样理解并遵循指令。

需要注意的是，OpenAI 最初是使用 TensorFlow 保存 GPT-2 的权重的，因此我们需要安装 TensorFlow 才能在 Python 中加载这些权重。此外，以下代码还会使用一个名为 tqdm 的进度条工具来跟踪下载进度，因此我们也需要安装它。

你可以在终端中执行以下命令来安装这些库：

pip3 install tensorflow>=2.15.0 tqdm>=4.66

from importlib.metadata import version
print("TensorFlow version:", version("tensorflow"))
print("tqdm version:", version("tqdm"))

TensorFlow version: 2.19.0
tqdm version: 4.67.1

下载代码相对较长，大多是样板代码，逻辑上也不复杂。因此，为了节省本章的宝贵篇幅，我们不再讨论从互联网获取文件的 Python 代码，而是直接从本章的在线仓库中下载 gpt_download.py 模块：

import urllib.request
url = (
    "https://raw.githubusercontent.com/rasbt/"
    "LLMs-from-scratch/main/ch05/"
    "01_main-chapter-code/gpt_download.py"
)
filename = url.split('/')[-1]
urllib.request.urlretrieve(url, filename)

接下来，在将该文件下载到 Python 会话的本地目录后，建议读者简单查看一下该文件的内容，以确保它被正确保存，并且包含有效的 Python 代码。

现在我们可以从 gpt_download.py 文件中导入 download_and_load_gpt2 函数，如下所示。该函数将 GPT-2 的架构设置（settings）和权重参数（params）加载到我们的 Python 会话中：

from gpt_download import download_and_load_gpt2
settings, params = download_and_load_gpt2(model_size="124M", models_dir="gpt2")

执行上述代码后，将下载与 124M 参数的 GPT-2 模型相关的以下 7 个文件：

checkpoint: 100%|██████████| 77.0/77.0 [00:00<00:00, 77.1kiB/s]
encoder.json: 100%|██████████| 1.04M/1.04M [00:01<00:00, 870kiB/s]
hparams.json: 100%|██████████| 90.0/90.0 [00:00<00:00, 36.2kiB/s]
model.ckpt.data-00000-of-00001: 100%|██████████| 498M/498M [06:26<00:00, 1.29MiB/s]  
model.ckpt.index: 100%|██████████| 5.21k/5.21k [00:00<00:00, 4.84MiB/s]
model.ckpt.meta: 100%|██████████| 471k/471k [00:00<00:00, 516kiB/s] 
vocab.bpe: 100%|██████████| 456k/456k [00:02<00:00, 213kiB/s] 

更新的下载说明
如果下载代码对你不起作用，可能是由于间歇性的网络连接问题、服务器故障，或者 OpenAI 共享开源 GPT-2 模型权重的方式发生了变化。在这种情况下，请访问本章的在线代码仓库：https://github.com/rasbt/LLMs-from-scratch 以获取替代和最新的说明。如有进一步问题，也请通过 Manning 论坛联系。

在上述代码执行完成之后，让我们检查 settings 和 params 的内容：

print("Settings:", settings)
print("Parameter dictionary keys:", params.keys())

输出内容如下：

Settings: {'n_vocab': 50257, 'n_ctx': 1024, 'n_embd': 768, 'n_head': 12, 'n_layer': 12}
Parameter dictionary keys: dict_keys(['blocks', 'b', 'g', 'wpe', 'wte'])

settings 和 params 都是 Python 字典。settings 字典保存了 LLM 的架构设置，类似于我们手动定义的 GPT_CONFIG_124M 设置。而 params 字典包含了实际的权重张量。请注意，我们这里只打印了字典的键名，因为直接打印权重内容会占据过多的屏幕空间。不过我们可以通过打印整个字典 print(params) 或者通过相应的字典键名来选择并检查各个权重张量，例如嵌入层的权重：

print(params["wte"])
print("Token embedding weight tensor dimensions:", params["wte"].shape)

标记嵌入层的权重如下：

[[-0.11010301 -0.03926672  0.03310751 ... -0.1363697   0.01506208
   0.04531523]
 [ 0.04034033 -0.04861503  0.04624869 ...  0.08605453  0.00253983
   0.04318958]
 [-0.12746179  0.04793796  0.18410145 ...  0.08991534 -0.12972379
  -0.08785918]
 ...
 [-0.04453601 -0.05483596  0.01225674 ...  0.10435229  0.09783269
  -0.06952604]
 [ 0.1860082   0.01665728  0.04611587 ... -0.09625227  0.07847701
  -0.02245961]
 [ 0.05135201 -0.02768905  0.0499369  ...  0.00704835  0.15519823
   0.12067825]]
Token embedding weight tensor dimensions: (50257, 768)

我们通过 download_and_load_gpt2(model_size="124M", ...) 设置下载并加载了最小的 GPT-2 模型的权重。然而，请注意，OpenAI 也分享了更大模型的权重：“355M”、“774M” 和 “1558M”。这些不同大小的 GPT 模型整体架构是相同的，如图 5.17 所示。

图 5.17 GPT-2 LLM 有多个不同的模型规模，从 1.24 亿参数到 15.58 亿参数不等。核心架构是相同的，唯一的区别在于嵌入维度大小以及注意力头和 transformer 块等组件重复的次数。

如图 5.17 所示，不同大小的 GPT-2 模型的整体架构保持相同，区别仅在于某些架构组件被重复的次数不同，嵌入向量的维度也有所不同。本章剩余的代码同样适用于这些更大的模型。

在将 GPT-2 模型权重加载到 Python 中之后，我们仍需将它们从 settings 和 params 字典中转移到我们的 GPTModel 实例中。

首先，我们创建一个字典，列出不同 GPT 模型规模之间的差异，如图 5.17 所示：

model_configs = {
  "gpt2-small (124M)": {"emb_dim": 768, "n_layers": 12, "n_heads": 12},
  "gpt2-medium (355M)": {"emb_dim": 1024, "n_layers": 24, "n_heads": 16},
  "gpt2-large (774M)": {"emb_dim": 1280, "n_layers": 36, "n_heads": 20},
  "gpt2-xl (1558M)": {"emb_dim": 1600, "n_layers": 48, "n_heads": 25},
}

假设我们希望加载最小的模型 "gpt2-small (124M)"，我们可以使用 model_configs 表中对应的设置来更新我们在本章前面使用的 GPT_CONFIG_124M 配置字典：

# Copy the base configuration and update with specific model settings
model_name = "gpt2-small (124M)"  # Example model name
NEW_CONFIG = GPT_CONFIG_124M.copy()
NEW_CONFIG.update(model_configs[model_name])

细心的读者可能还记得，我们之前使用的是 256-token 的长度，而 OpenAI 的原始 GPT-2 模型是使用 1024-token 长度进行训练的，因此我们需要相应地更新 NEW_CONFIG：

NEW_CONFIG.update({"context_length": 1024})

此外，OpenAI 在多头注意力模块的线性层中使用了偏置向量（bias）来实现 query、key 和 value 的矩阵计算。尽管偏置向量在当前的 LLM 中已不再常用，因为它们对建模性能没有明显提升，但由于我们现在使用的是预训练权重，因此需要启用这些偏置向量以保持设置一致性：

NEW_CONFIG.update({"qkv_bias": True})

现在我们可以使用更新后的 NEW_CONFIG 字典初始化一个新的 GPTModel 实例：

gpt = GPTModel(NEW_CONFIG)
gpt.eval()

默认情况下，GPTModel 实例会使用随机权重进行初始化以便于预训练。要使用 OpenAI 的模型权重，最后一步就是将这些随机权重用我们在 params 字典中加载的权重进行替换。

为此，我们首先定义一个小型的 assign 工具函数，用于检查两个张量或数组（left 和 right）是否具有相同的维度或形状，并将右侧的张量作为可训练的 PyTorch 参数返回：

def assign(left, right):
    if left.shape != right.shape:
        raise ValueError(f"Shape mismatch. Left: {left.shape}, Right: {right.shape}")
    return torch.nn.Parameter(torch.tensor(right))

接下来，我们定义一个 load_weights_into_gpt 函数，用于将 params 字典中的权重加载到 GPTModel 实例 gpt 中：

代码清单 5.5 将 OpenAI 权重加载进我们的 GPT 模型代码中

import numpy as np

def load_weights_into_gpt(gpt, params):
    gpt.pos_emb.weight = assign(gpt.pos_emb.weight, params['wpe'])
    gpt.tok_emb.weight = assign(gpt.tok_emb.weight, params['wte'])
    
    for b in range(len(params["blocks"])):
        q_w, k_w, v_w = np.split(
            (params["blocks"][b]["attn"]["c_attn"])["w"], 3, axis=-1)
        gpt.trf_blocks[b].att.W_query.weight = assign(
            gpt.trf_blocks[b].att.W_query.weight, q_w.T)
        gpt.trf_blocks[b].att.W_key.weight = assign(
            gpt.trf_blocks[b].att.W_key.weight, k_w.T)
        gpt.trf_blocks[b].att.W_value.weight = assign(
            gpt.trf_blocks[b].att.W_value.weight, v_w.T)

        q_b, k_b, v_b = np.split(
            (params["blocks"][b]["attn"]["c_attn"])["b"], 3, axis=-1)
        gpt.trf_blocks[b].att.W_query.bias = assign(
            gpt.trf_blocks[b].att.W_query.bias, q_b)
        gpt.trf_blocks[b].att.W_key.bias = assign(
            gpt.trf_blocks[b].att.W_key.bias, k_b)
        gpt.trf_blocks[b].att.W_value.bias = assign(
            gpt.trf_blocks[b].att.W_value.bias, v_b)

        gpt.trf_blocks[b].att.out_proj.weight = assign(
            gpt.trf_blocks[b].att.out_proj.weight, 
            params["blocks"][b]["attn"]["c_proj"]["w"].T)
        gpt.trf_blocks[b].att.out_proj.bias = assign(
            gpt.trf_blocks[b].att.out_proj.bias, 
            params["blocks"][b]["attn"]["c_proj"]["b"])

        gpt.trf_blocks[b].ff.layers[0].weight = assign(
            gpt.trf_blocks[b].ff.layers[0].weight, 
            params["blocks"][b]["mlp"]["c_fc"]["w"].T)
        gpt.trf_blocks[b].ff.layers[0].bias = assign(
            gpt.trf_blocks[b].ff.layers[0].bias, 
            params["blocks"][b]["mlp"]["c_fc"]["b"])
        gpt.trf_blocks[b].ff.layers[2].weight = assign(
            gpt.trf_blocks[b].ff.layers[2].weight, 
            params["blocks"][b]["mlp"]["c_proj"]["w"].T)
        gpt.trf_blocks[b].ff.layers[2].bias = assign(
            gpt.trf_blocks[b].ff.layers[2].bias, 
            params["blocks"][b]["mlp"]["c_proj"]["b"])

        gpt.trf_blocks[b].norm1.scale = assign(
            gpt.trf_blocks[b].norm1.scale, 
            params["blocks"][b]["ln_1"]["g"])
        gpt.trf_blocks[b].norm1.shift = assign(
            gpt.trf_blocks[b].norm1.shift, 
            params["blocks"][b]["ln_1"]["b"])
        gpt.trf_blocks[b].norm2.scale = assign(
            gpt.trf_blocks[b].norm2.scale, 
            params["blocks"][b]["ln_2"]["g"])
        gpt.trf_blocks[b].norm2.shift = assign(
            gpt.trf_blocks[b].norm2.shift, 
            params["blocks"][b]["ln_2"]["b"])

    gpt.final_norm.scale = assign(gpt.final_norm.scale, params["g"])
    gpt.final_norm.shift = assign(gpt.final_norm.shift, params["b"])
    gpt.out_head.weight = assign(gpt.out_head.weight, params["wte"])
    
    
load_weights_into_gpt(gpt, params)
gpt.to(device)

在 load_weights_into_gpt 函数中，我们小心地将 OpenAI 实现中的权重与我们 GPTModel 实现中的权重一一对应起来。举一个具体的例子，OpenAI 将第一个 Transformer 模块中输出投影层的权重张量存储为 params["blocks"][0]["attn"]["c_proj"]["w"]。在我们的实现中，该权重张量对应于 gpt.trf_blocks[b].att.out_proj.weight，其中 gpt 是一个 GPTModel 实例。

开发 load_weights_into_gpt 函数的过程涉及大量的猜测，因为 OpenAI 使用的命名方式与我们的略有不同。然而，如果尝试匹配两个形状不一致的张量，assign 函数会提醒我们出错。此外，如果我们在该函数中犯了错误，我们会注意到，因为生成的 GPT 模型将无法产生连贯的文本。

现在我们来实际尝试运行 load_weights_into_gpt 并将 OpenAI 的模型权重加载到我们的 GPTModel 实例 gpt 中：

load_weights_into_gpt(gpt, params)
gpt.to(device)

如果模型加载正确，我们现在就可以使用我们之前定义的 generate 函数生成新文本：

torch.manual_seed(123)

token_ids = generate(
    model=gpt,
    idx=text_to_token_ids("Every effort moves you", tokenizer).to(device),
    max_new_tokens=25,
    context_size=NEW_CONFIG["context_length"],
    top_k=50,
    temperature=1.5
)

print("Output text:\n", token_ids_to_text(token_ids, tokenizer))

生成的文本如下：

Output text:
 Every effort moves you as far as the hand can go until the end of your turn unless something interrupts your control flow. As you may observe I

我们可以确信我们正确加载了模型权重，因为模型可以生成连贯的文本。如果在这个过程中出现哪怕是极小的错误，模型也将无法正常工作。

在接下来的章节中，我们将继续使用这个预训练模型，并对其进行微调，使其能够进行文本分类和遵循指令。

练习 5.5
计算加载了 OpenAI 预训练权重的 GPTModel 在 “The Verdict” 数据集上的训练集和验证集损失。

import tiktoken
import torch
from previous_chapters import GPTModel


GPT_CONFIG_124M = {
    "vocab_size": 50257,   # Vocabulary size
    "context_length": 256, # Shortened context length (orig: 1024)
    "emb_dim": 768,        # Embedding dimension
    "n_heads": 12,         # Number of attention heads
    "n_layers": 12,        # Number of layers
    "drop_rate": 0.1,      # Dropout rate
    "qkv_bias": False      # Query-key-value bias
}


torch.manual_seed(123)

tokenizer = tiktoken.get_encoding("gpt2")

from gpt_download import download_and_load_gpt2

settings, params = download_and_load_gpt2(model_size="124M", models_dir="gpt2")

# Define model configurations in a dictionary for compactness
model_configs = {
    "gpt2-small (124M)": {"emb_dim": 768, "n_layers": 12, "n_heads": 12},
    "gpt2-medium (355M)": {"emb_dim": 1024, "n_layers": 24, "n_heads": 16},
    "gpt2-large (774M)": {"emb_dim": 1280, "n_layers": 36, "n_heads": 20},
    "gpt2-xl (1558M)": {"emb_dim": 1600, "n_layers": 48, "n_heads": 25},
}

# Copy the base configuration and update with specific model settings
model_name = "gpt2-small (124M)"  # Example model name
NEW_CONFIG = GPT_CONFIG_124M.copy()
NEW_CONFIG.update(model_configs[model_name])
NEW_CONFIG.update({"context_length": 1024, "qkv_bias": True})

gpt = GPTModel(NEW_CONFIG)
gpt.eval();

from gpt_generate import load_weights_into_gpt


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
load_weights_into_gpt(gpt, params)
gpt.to(device);
import os
import urllib.request
from previous_chapters import create_dataloader_v1


file_path = "the-verdict.txt"
url = "https://raw.githubusercontent.com/rasbt/LLMs-from-scratch/main/ch02/01_main-chapter-code/the-verdict.txt"

if not os.path.exists(file_path):
    with urllib.request.urlopen(url) as response:
        text_data = response.read().decode('utf-8')
    with open(file_path, "w", encoding="utf-8") as file:
        file.write(text_data)
else:
    with open(file_path, "r", encoding="utf-8") as file:
        text_data = file.read()


# Train/validation ratio
train_ratio = 0.90
split_idx = int(train_ratio * len(text_data))
train_data = text_data[:split_idx]
val_data = text_data[split_idx:]


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
)
from gpt_train import calc_loss_loader

torch.manual_seed(123) # For reproducibility due to the shuffling in the data loader
train_loss = calc_loss_loader(train_loader, gpt, device)
val_loss = calc_loss_loader(val_loader, gpt, device)

print("Training loss:", train_loss)
print("Validation loss:", val_loss)

gpt_train.py

Training loss: 3.7547486888037787
Validation loss: 3.5596182346343994

settings, params = download_and_load_gpt2(model_size="1558M", models_dir="gpt2")

model_name = "gpt2-xl (1558M)"
NEW_CONFIG = GPT_CONFIG_124M.copy()
NEW_CONFIG.update(model_configs[model_name])
NEW_CONFIG.update({"context_length": 1024, "qkv_bias": True})

gpt = GPTModel(NEW_CONFIG)
gpt.eval()

load_weights_into_gpt(gpt, params)
gpt.to(device)

torch.manual_seed(123)
train_loss = calc_loss_loader(train_loader, gpt, device)
val_loss = calc_loss_loader(val_loader, gpt, device)

print("Training loss:", train_loss)
print("Validation loss:", val_loss)

Training loss: 3.3046312861972384
Validation loss: 3.1195147037506104

练习 5.6
鼓励读者尝试使用不同大小的 GPT-2 模型，例如最大的 1558M 参数模型，并将其生成的文本与本章中加载的 124M 模型进行比较。

import tiktoken
import torch
from previous_chapters import GPTModel


GPT_CONFIG_124M = {
    "vocab_size": 50257,   # Vocabulary size
    "context_length": 256, # Shortened context length (orig: 1024)
    "emb_dim": 768,        # Embedding dimension
    "n_heads": 12,         # Number of attention heads
    "n_layers": 12,        # Number of layers
    "drop_rate": 0.1,      # Dropout rate
    "qkv_bias": False      # Query-key-value bias
}


tokenizer = tiktoken.get_encoding("gpt2")
from gpt_download import download_and_load_gpt2
from gpt_generate import load_weights_into_gpt


model_configs = {
    "gpt2-small (124M)": {"emb_dim": 768, "n_layers": 12, "n_heads": 12},
    "gpt2-medium (355M)": {"emb_dim": 1024, "n_layers": 24, "n_heads": 16},
    "gpt2-large (774M)": {"emb_dim": 1280, "n_layers": 36, "n_heads": 20},
    "gpt2-xl (1558M)": {"emb_dim": 1600, "n_layers": 48, "n_heads": 25},
}

model_name = "gpt2-xl (1558M)"
NEW_CONFIG = GPT_CONFIG_124M.copy()
NEW_CONFIG.update(model_configs[model_name])
NEW_CONFIG.update({"context_length": 1024, "qkv_bias": True})

gpt = GPTModel(NEW_CONFIG)
gpt.eval()

settings, params = download_and_load_gpt2(model_size="1558M", models_dir="gpt2")
print(load_weights_into_gpt(gpt, params))

File already exists and is up-to-date: gpt2/1558M/checkpoint
File already exists and is up-to-date: gpt2/1558M/encoder.json
File already exists and is up-to-date: gpt2/1558M/hparams.json
File already exists and is up-to-date: gpt2/1558M/model.ckpt.data-00000-of-00001
File already exists and is up-to-date: gpt2/1558M/model.ckpt.index
File already exists and is up-to-date: gpt2/1558M/model.ckpt.meta
File already exists and is up-to-date: gpt2/1558M/vocab.bpe

from gpt_generate import generate, text_to_token_ids, token_ids_to_text
torch.manual_seed(123)

token_ids = generate(
    model=gpt,
    idx=text_to_token_ids("Every effort moves you", tokenizer),
    max_new_tokens=25,
    context_size=NEW_CONFIG["context_length"],
    top_k=50,
    temperature=1.5
)

print("Output text:\n", token_ids_to_text(token_ids, tokenizer))

Output text:
 Every effort moves you toward finding an ideal life. You don't have to accept your current one at once, because if you do you'll never

5.6 小结

• 当LLM生成文本时，它一次输出一个token。
• 默认情况下，下一个token的生成方式是将模型输出转换为概率得分，并从词汇表中选择对应于最大概率得分的token，这种方式称为“贪婪解码”。
• 通过使用概率采样和温度缩放，我们可以影响生成文本的多样性和连贯性。
• 训练集和验证集的损失可以用于评估LLM在训练过程中生成文本的质量。
• 预训练LLM的过程涉及调整其权重以最小化训练损失。
• LLM的训练循环本身是深度学习中的标准流程，使用常规的交叉熵损失函数和AdamW优化器。
• 在大型文本语料库上预训练LLM耗时且资源密集，因此我们可以加载OpenAI公开提供的权重，作为我们自己在大数据集上预训练模型的替代方案。