Build a Large Language Model (From Scratch) · Chapter 5 中文译文
在无标签数据上预训练
译自 Build a Large Language Model (From Scratch) MEAP V08。本文保留章节结构、代码、图、表与公式标记；图像从原 PDF 页段抽取并嵌入。
第 5 章 在无标签数据上预训练
本章涵盖以下内容
- 计算训练集和验证集损失，以评估训练期间 LLM 生成文本的质量
- 实现训练函数并预训练 LLM
- 保存和加载模型权重，以继续训练 LLM
- 加载来自 OpenAI 的预训练权重
在前几章中，我们实现了数据采样、注意力机制，并编写了 LLM 架构。本章的核心重点是实现训练函数并预训练 LLM，如图 5.1 所示。
如图 5.1 所示，我们还将学习基本的模型评估技术，用于衡量生成文本的质量；这是在训练过程中优化 LLM 的必要条件。此外，我们还会讨论如何加载预训练权重，从而为后续章节中的微调提供一个可靠的起点。
5.1 评估生成式文本模型
本章一开始，我们会基于上一章的代码设置用于文本生成的 LLM，并在本节讨论评估生成文本质量的基本方法。本节以及本章其余部分涵盖的内容如图 5.2 所示。
如图 5.2 所示，在进入后续小节中的文本评估以及训练损失和验证损失计算之前，下一小节会先回顾我们在上一章末尾搭建的文本生成过程。
5.1.1 使用 GPT 生成文本
在本节中，我们会设置 LLM，并简要回顾第 4 章实现的文本生成过程。我们先使用第 4 章中的 GPTModel 类和 GPT_CONFIG_124M 字典，初始化本章将要评估和训练的 GPT 模型：
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()#A我们将上下文长度从 \(1024\) 个 token 缩短到 \(256\) 个 token#B将 dropout 设置为 \(0\) 是可行且常见的做法。
考虑 GPT_CONFIG_124M 字典，与上一章相比，我们唯一做出的调整是把上下文长度（context_length）减少到 \(256\) 个 token。这个修改降低了训练模型的计算需求，使得在标准笔记本电脑上执行训练成为可能。
最初，拥有 \(124\) 百万个参数的 GPT-2 模型被配置为最多处理 \(1,024\) 个 token。在本章末尾完成训练过程之后，我们会更新上下文大小设置，并加载预训练权重，以使用配置为 \(1,024\)-token 上下文长度的模型。
使用 GPTModel 实例后，我们采用上一章介绍的 generate_text_simple 函数，并引入两个方便的函数：text_to_token_ids 和 token_ids_to_text。这些函数便于在文本表示和 token 表示之间进行转换，这项技术将在本章反复使用。为了更清楚地理解，在进入代码之前，图 5.3 展示了这个过程。
图 5.3 展示了使用 GPT 模型进行文本生成的三步流程。第一，tokenizer 将输入文本转换成一系列 token ID，如第 2 章所讨论。第二，模型接收这些 token ID 并生成对应的 logits；logits 是表示词汇表中每个 token 概率分布的向量，如第 4 章所讨论。第三，这些 logits 被转换回 token ID，tokenizer 再将其解码成人类可读的文本，由此完成从文本输入到文本输出的循环。
在代码中，我们按如下方式实现文本生成过程：
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) # add batch dimension
return encoded_tensor
def token_ids_to_text(token_ids, tokenizer):
flat = token_ids.squeeze(0) # remove batch dimension
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 计算文本生成损失
本节会探索如何通过计算所谓的文本生成损失，对训练期间生成的文本质量进行数值评估。为了让概念清楚且可应用，我们会结合一个实践示例逐步讲解这个主题；首先简要回顾第 2 章中数据如何加载，以及第 4 章中的 generate_text_simple 函数如何生成文本。
图 5.4 展示了从输入文本到 LLM 生成文本的整体流程，其中包含五个步骤。
图 5.4 中的文本生成过程概述了第 4 章的 generate_text_simple 函数在内部执行的操作。在本节稍后计算用于衡量生成文本质量的损失之前，我们需要先执行这些相同的初始步骤。
图 5.4 使用一个小型 \(7\)-token 词汇表来展示文本生成过程，以便让这张图放在单页中。不过，我们的 GPTModel 使用一个大得多的词汇表，其中包含 \(50,257\) 个词；因此，下面代码中的 token ID 范围将是 \(0\) 到 \(50,256\)，而不是 \(0\) 到 \(6\)。
此外，为了简单起见，图 5.4 只展示了一个文本示例（“every effort moves”）。在下面实现图 5.4 步骤的动手代码示例中，我们将使用两个输入示例（“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",
[107, 588, 11311]]) # " really like chocolate"]请注意，目标是输入向前移动一个位置之后的结果；这个概念我们在第 2 章实现数据加载器时介绍过。这种移位策略对于教会模型预测序列中的下一个 token 至关重要。
当我们把输入送入模型，为两个输入示例计算 logit 向量时，每个示例都包含三个 token；随后应用 softmax 函数把这些 logit 值转换成概率分数，这对应图 5.4 中的第 2 步：
with torch.no_grad(): #A
logits = model(inputs)
probas = torch.softmax(logits, dim=-1) # Probability of each token in vocabulary
print(probas.shape)#A由于我们还没有进行训练，因此禁用梯度跟踪
得到的概率分数张量（probas）的张量维度如下：
torch.Size([2, 3, 50257])第一个数字 \(2\) 对应输入中的两个示例（行），也称为批次大小。第二个数字 \(3\) 对应每个输入（行）中的 token 数量。最后一个数字对应嵌入维度，它由词汇表大小决定，如前几章所讨论。
通过 softmax 函数把 logits 转换成概率之后，第 4 章中的 generate_text_simple 函数会将得到的概率分数转换回文本，如图 5.4 的第 3 到第 5 步所示。
我们可以通过对概率分数应用 argmax 函数来实现第 3 步和第 4 步，从而得到对应的 token ID：
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], # First batch
[ 339],
[42826]],
[[49906], # Second batch
[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 与我们希望模型生成的目标 token 相差很大：
Targets batch 1: effort moves you
Outputs batch 1: Armed heNetflix模型生成了与目标文本不同的随机文本，因为它尚未经过训练。现在我们进入通过所谓损失对模型生成文本的性能进行数值评估的部分，如图 5.4 所示。这不仅有助于衡量生成文本的质量，也是稍后实现训练函数的构建块；我们会使用训练函数更新模型权重，以改善生成文本。
如图 5.5 所示，本节余下部分要实现的文本评估过程，其中一部分是衡量生成的 token 与正确预测（目标）之间“相距多远”。本章稍后实现的训练函数将使用这些信息来调整模型权重，使生成的文本更接近（理想情况下匹配）目标文本。
模型训练的目标是提高与正确目标 token ID 对应的索引位置上的 softmax 概率，如图 5.6 所示。这个 softmax 概率也用于本节余下部分实现的评估指标，用来对模型生成输出进行数值评估：正确位置上的概率越高，效果越好。
请记住，图 5.6 显示的是一个紧凑 \(7\)-token 词汇表上的 softmax 概率，以便把所有内容放进一张图中。这意味着初始随机值会在 \(1/7\) 附近，也就是约 \(0.14\)。
然而，我们为 GPT-2 模型使用的词汇表有 \(50,257\) 个 token，因此大多数初始概率会通过 \(1/50,257\) 落在 \(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)每个批次的 \(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。
在本节余下部分，我们会为两个示例批次 target_probas_1 和 target_probas_2 的概率分数计算损失。主要步骤如图 5.7 所示。
由于我们已经应用图 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])在数学优化中，处理概率分数的对数比直接处理分数更容易。本主题超出了本书范围，但我在一节讲座中做了进一步说明，该讲座链接在附录 B 的参考资料部分。
接下来，我们通过计算平均值把这些对数概率合并为单个分数（图 5.7 中的第 5 步）：
avg_log_probas = torch.mean(log_probas)
print(avg_log_probas)得到的平均对数概率分数如下：
tensor(-10.7940)目标是在训练过程中更新模型权重，使平均对数概率尽可能接近 \(0\)；我们将在第 5.2 节实现这个训练过程。
不过，在深度学习中，常见做法不是把平均对数概率推高到 \(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 中的全部 \(6\) 个步骤。
在应用交叉熵函数之前，我们先简要回顾 logits 和目标张量的形状：
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 数量和词汇表大小。目标张量有两个维度：批次大小和 token 数量。
对于 PyTorch 中的 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])请记住，目标是我们希望 LLM 生成的 token ID，而 logits 包含模型进入 softmax 函数以获得概率分数之前的未缩放输出。
之前，我们应用 softmax 函数，选择与目标 ID 对应的概率分数，并计算负平均对数概率。PyTorch 的 cross_entropy 函数会替我们处理所有这些步骤：
loss = torch.nn.functional.cross_entropy(logits_flat, targets_flat)
print(loss)得到的损失与我们之前手动应用图 5.7 中各个步骤时得到的结果相同：
tensor(10.7940)在本节中，我们为了说明目的计算了两个小型文本输入的损失。下一节中，我们会把损失计算应用到整个训练集和验证集。
5.1.3 计算训练集和验证集损失
在本节中，我们首先准备稍后在本章用于训练 LLM 的训练数据集和验证数据集。然后，我们计算训练集和验证集的交叉熵，如图 5.8 所示；这是模型训练过程中的一个重要组件。
为了计算图 5.8 所示训练数据集和验证数据集上的损失，我们使用一个非常小的文本数据集，即 Edith Wharton 的短篇小说 “The Verdict”；我们已经在第 2 章使用过它。通过选择公有领域文本，我们避免了与使用权相关的任何顾虑。此外，我们之所以使用如此小的数据集，是因为它可以让代码示例在标准笔记本电脑上运行，只需几分钟，即使没有高端 GPU 也是如此；这对教学目的尤其有利。
感兴趣的读者也可以使用本书的补充代码，准备一个更大规模的数据集，其中包含来自 Project Gutenberg 的 \(60,000\) 多本公有领域图书，并在这些图书上训练 LLM（详情参见附录 D）。
下面的代码加载我们在第 2 章使用过的 “The Verdict” 短篇小说：
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 章的数据加载器来准备 LLM 训练用的批次。这个过程在图 5.9 中可视化展示。
出于可视化目的，图 5.9 由于空间限制使用了 max_length=6。不过，对于我们实际实现的数据加载器，我们会把 max_length 设置为等于 LLM 支持的 \(256\)-token 上下文长度，使 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 子集后，我们现在可以复用第 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
)在前面的代码中，我们使用了相对较小的批次大小，以降低计算资源需求，因为我们使用的是非常小的数据集。在实践中，用 \(1,024\) 或更大批次大小训练 LLM 并不少见。
作为一个可选检查，我们可以遍历数据加载器，以确保它们被正确创建：
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\) 个训练集批次，每个批次包含 \(2\) 个样本，每个样本有 \(256\) 个 token。由于我们只分配了 \(10\%\) 的数据用于验证，因此只有一个验证批次，其中包含 \(2\) 个输入示例。
如预期一样，输入数据（x）和目标数据（y）具有相同形状（批次大小乘以每个批次中的 token 数量），因为目标就是输入向前移动一个位置后的结果，如第 2 章所讨论。
接下来，我们实现一个实用函数，用于计算训练加载器和验证加载器返回的给定批次的交叉熵损失：
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#A转移到给定设备，使我们能够把数据转移到 GPU
现在，我们可以使用这个为单个批次计算损失的 calc_loss_batch 实用函数，来实现下面的 calc_loss_loader 函数；它会计算给定数据加载器采样到的所有批次上的损失：
def calc_loss_loader(data_loader, model, device, num_batches=None):
total_loss = 0.
if len(data_loader) == 0:
return float("nan")
elif 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#A如果没有指定固定的num_batches，则遍历所有批次#B如果num_batches超过数据加载器中的批次数量，则减少批次数量以匹配数据加载器中的总批次数量#C对每个批次的损失求和#D对所有批次的损失取平均
默认情况下，calc_loss_batch 函数会遍历给定数据加载器中的所有批次，把损失累加到 total_loss 变量中，然后根据总批次数计算并平均损失。或者，我们可以通过 num_batches 指定较小的批次数量，以便在模型训练期间加快评估速度。
现在让我们看看这个 calc_loss_batch 函数的实际效果，把它应用到训练集加载器和验证集加载器：
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") #A
model.to(device)
with torch.no_grad(): #B
train_loss = calc_loss_loader(train_loader, model, device) #C
val_loss = calc_loss_loader(val_loader, model, device)
print("Training loss:", train_loss)
print("Validation loss:", val_loss)#A如果你的机器有支持 CUDA 的 GPU，那么无需修改任何代码，LLM 就会在 GPU 上训练#B由于我们还没有进行训练，因此为了效率禁用梯度跟踪#C通过device设置，我们确保数据被加载到与 LLM 模型相同的设备上
得到的损失值如下：
Training loss: 10.98758347829183
Validation loss: 10.98110580444336损失值相对较高，因为模型尚未经过训练。作为对比，如果模型学会按训练集和验证集中出现的方式生成下一个 token，那么损失会接近 \(0\)。
现在，我们已经有办法衡量生成文本的质量；下一节中，我们会训练 LLM 来降低这个损失，使它更擅长生成文本，如图 5.10 所示。
如图 5.10 所示，下一节将重点关注预训练 LLM。模型训练完成后，我们会实现其他文本生成策略，并保存和加载预训练模型权重。
5.2 训练 LLM
在本节中，我们终于要实现用于预训练 LLM 的代码，也就是预训练我们的 GPTModel。为此，我们聚焦于一个简单直接的训练循环，如图 5.11 所示，以保持代码简洁且易读。不过，感兴趣的读者可以在附录 D《为训练循环添加各种增强》中了解更高级的技术，包括学习率预热、余弦退火和梯度裁剪。
图 5.11 中的流程图展示了一个典型的 PyTorch 神经网络训练工作流，我们将用它来训练 LLM。该流程概括了八个步骤，从遍历每个 epoch 开始，处理批次，重置并计算梯度，更新权重，最后执行监控步骤，例如打印损失和生成文本样本。如果你对使用 PyTorch 训练深度神经网络还比较陌生，并且对其中任何步骤不熟悉，可以阅读附录 A《PyTorch 简介》中的 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):
train_losses, val_losses, track_tokens_seen = [], [], [] #A
tokens_seen, global_step = 0, -1
for epoch in range(num_epochs): #B
model.train()
for input_batch, target_batch in train_loader:
optimizer.zero_grad() #C
loss = calc_loss_batch(input_batch, target_batch, model, device)
loss.backward() #D
optimizer.step() #E
tokens_seen += input_batch.numel()
global_step += 1
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}")
generate_and_print_sample( #G
model, tokenizer, device, start_context
)
return train_losses, val_losses, track_tokens_seen#A初始化列表，用于跟踪损失和已见过的 token 数量#B开始主训练循环#C重置前一个批次迭代中的损失梯度#D计算损失梯度#E使用损失梯度更新模型权重#F可选的评估步骤#G在每个 epoch 后打印一个样本文本
请注意，我们刚刚创建的 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#A在评估期间禁用 dropout，以获得稳定且可复现的结果#B禁用梯度跟踪；评估时不需要它，这样可以降低计算开销
与 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):
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()evaluate_model 函数为我们提供模型训练进展的数值估计，而这个 generate_and_print_sample 文本函数则提供一个由模型生成的具体文本示例，用于在训练过程中判断模型的能力。
现在，让我们把这些都放到实践中：使用 AdamW 优化器和前面定义的 train_model_simple 函数，训练一个 GPTModel 实例 \(10\) 个 epoch。
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) #A
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=1,
start_context="Every effort moves you", tokenizer=tokenizer
)#A.parameters()方法返回模型的所有可训练权重参数
执行 training_model_simple 函数会启动训练过程；在 MacBook Air 或类似笔记本电脑上，完成该过程大约需要 \(5\) 分钟。执行期间打印的输出如下：
Ep 1 (Step 000000): Train loss 9.781, Val loss 9.933
Ep 1 (Step 000005): Train loss 8.111, Val loss 8.339
Every effort moves you,,,,,,,,,,,,.
Ep 2 (Step 000010): Train loss 6.661, Val loss 7.048
Ep 2 (Step 000015): Train loss 5.961, Val loss 6.616
Every effort moves you, and, and, and, and, and, and, and, and, and, and, and, and, and,
and, and, and, and, and, and, and, and, and,, and, and,
[...] Results are truncated to save space #A
Ep 9 (Step 000080): Train loss 0.541, Val loss 6.393
Every effort moves you?" "Yes--quite insensible to the irony. She wanted him
vindicated--and by me!" He laughed again, and threw back the window-curtains, I had the
donkey. "There were days when I
Ep 10 (Step 000085): Train loss 0.391, Val loss 6.452
Every effort moves you know," was one of the axioms he laid down across the Sevres and
silver of an exquisitely appointed luncheon-table, when, on a later day, I had again run
over from Monte Carlo; and Mrs. Gis#A为节省空间，移除了中间结果
可以看到，根据训练期间打印出的结果，训练损失大幅改善，从 \(9.558\) 开始并收敛到 \(0.762\)。模型的语言能力也有了很大提升。起初，模型只能在起始上下文后追加逗号（“Every effort moves you,,,,,,,,,,,,”）或重复单词 “and”。到训练结束时，它已经能够生成语法正确的文本。
与训练集损失类似，我们可以看到验证损失一开始也很高（\(9.856\)），并且在训练期间下降。不过，它始终没有变得像训练集损失那么小，并且在第 \(10\) 个 epoch 后仍保持在 \(6.372\)。
在更详细地讨论验证损失之前，让我们先创建一个简单图表，将训练集损失和验证集损失并排显示：
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() #A
ax2.plot(tokens_seen, train_losses, alpha=0) #B
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)#A创建第二个 x 轴，并与同一个 y 轴共享#B用不可见的图来对齐刻度
得到的训练损失和验证损失图如图 5.12 所示。
如图 5.12 所示，在第一个 epoch 中，训练损失和验证损失都开始改善。不过，在第二个 epoch 之后，这两种损失开始分化。这种分化，以及验证损失远大于训练损失这一事实，都表明模型正在对训练数据过拟合。我们可以通过在 “The Verdict” 文本文件中搜索生成的文本片段（例如 “quite insensible to the irony”）来确认，模型确实逐字记住了训练数据。
这种记忆现象是预期之内的，因为我们使用的是一个非常非常小的训练数据集，并且对模型训练了多个 epoch。通常，更常见的做法是在大得多的数据集上只训练模型一个 epoch。
如前所述，感兴趣的读者可以尝试在来自 Project Gutenberg 的 \(60,000\) 本公版图书上训练模型；在那里不会出现这种过拟合。详情见附录 B。
在接下来的小节中，如图 5.13 所示，我们将探索 LLM 使用的采样方法，以减轻记忆效应，从而生成更新颖的文本。
如图 5.13 所示，下一节将介绍 LLM 的文本生成策略，用于减少对训练数据的记忆并提高 LLM 生成文本的原创性；之后我们会讨论权重的加载与保存，以及从 OpenAI 的 GPT 模型加载预训练权重。
5.3 用于控制随机性的解码策略
在本节中，我们将介绍文本生成策略（也称为解码策略），用于生成更具原创性的文本。首先，我们会简要回顾上一章中的 generate_text_simple 函数，也就是本章前面在 generate_and_print_sample 中使用过的函数。然后，我们将介绍两种技术：温度缩放和 top-k 采样，用来改进这个函数。
我们首先把模型从 GPU 转回 CPU，因为对于一个相对较小的模型，推理并不需要 GPU。此外，在训练之后，我们会把模型切换到评估模式，以关闭 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 know," was one of the axioms he laid down across the Sevres and
silver of an exquisitely appointed lun如第 5.1.2 节前面所解释的，在每个生成步骤中，生成的 token 都是从词表中所有 token 里选择概率分数最大的那个。
下面几个小节会介绍两个概念，用于控制生成文本的随机性和多样性：温度缩放和 top-k 采样。
5.3.1 温度缩放
本节介绍温度缩放，这是一种向下一个 token 生成任务加入概率选择过程的技术。
之前，在 generate_text_simple 函数内部，我们总是使用 torch.argmax 将概率最高的 token 采样为下一个 token，这也称为贪婪解码。为了生成更有变化的文本，我们可以用一个从概率分布中采样的函数来替换 argmax；这里的概率分布指的是 LLM 在每个 token 生成步骤中，为词表中每个条目生成的概率分数。
为了用一个具体示例说明概率采样，我们先用一个非常小的词表，简要讨论下一个 token 的生成过程：
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 对应的 token ID，然后可以通过反向词表把它映射回文本：
probas = torch.softmax(next_token_logits, dim=0)
next_token_id = torch.argmax(probas).item()
print(inverse_vocab[next_token_id])由于最大的 logit 值以及相应最大的 softmax 概率分数位于第 \(4\) 个位置（索引位置为 \(3\)，因为 Python 使用从 \(0\) 开始的索引），生成的词是 "forward"。
为了实现一个概率采样过程，我们现在可以用 PyTorch 中的 multinomial 函数替换 argmax：
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)
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]}")
print_sampled_tokens(probas)采样输出如下：
73 x closer
0 x every
0 x effort
582 x forward
2 x inches
0 x moves
0 x pizza
343 x toward从输出可以看出，词 "forward" 大多数时候会被采样到（\(1000\) 次中有 \(582\) 次），但其他 token，例如 "closer"、"inches" 和 "toward"，有时也会被采样到。这意味着，如果我们在 generate_and_print_sample 函数内部用 multinomial 函数替换 argmax 函数，LLM 有时会生成 "every effort moves you toward"、"every effort moves you inches" 和 "every effort moves you closer" 这样的文本，而不是 "every effort moves you forward"。
我们还可以通过一个叫作温度缩放的概念进一步控制分布和选择过程；温度缩放只是一个听起来更复杂的说法，本质上就是把 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] #A
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()
#A Original, lower, and higher confidence#A原始置信度、更低置信度和更高置信度
得到的图如图 5.14 所示。
"forward"）会具有更高的概率分数。反过来，把温度提高到 \(5\) 会让分布更均匀。温度为 \(1\) 时，会先将 logits 除以 \(1\)，再把它们传给 softmax 函数来计算概率分数。换句话说，使用温度 \(1\) 与不使用任何温度缩放是相同的。在这种情况下，token 会通过 PyTorch 中的 multinomial 采样函数，按照等于原始 softmax 概率分数的概率被选中。
另外，如图 5.14 所示，应用非常小的温度（例如 \(0.1\)）会得到更尖锐的分布，使得 multinomial 函数几乎 \(100\%\) 的时间都会选择最可能的 token（这里是 "forward"），从而接近 argmax 函数的行为。反过来，温度为 \(5\) 时会得到更均匀的分布，其他 token 会更频繁地被选中。这可以给生成文本增加更多变化，但也更常导致无意义的文本。例如，使用温度 \(5\) 时，大约 \(4\%\) 的时间会得到 "every effort moves you pizza" 这样的文本。
5.3.2 Top-k 采样
在上一节中，我们实现了一种结合温度缩放的概率采样方法，用来提高输出的多样性。我们看到，更高的温度值会带来更均匀分布的下一个 token 概率；由于这会降低模型反复选择最可能 token 的概率，因此会产生更多样的输出。这种方法允许在生成过程中探索可能性较低但潜在更有趣、更有创造性的路径。不过，这种方法的一个缺点是，它有时会导致语法错误或完全无意义的输出，例如 "every effort moves you pizza"。
在本节中，我们介绍另一个叫作 top-k 采样的概念；当它与概率采样和温度缩放结合使用时，可以改进文本生成结果。
在 top-k 采样中，我们可以把可采样的 token 限制为最可能的前 \(k\) 个 token，并通过掩蔽其他所有 token 的概率分数，把它们从选择过程中排除出去，如图 5.15 所示。
图 5.15 概述的方法会把所有未被选中的 logits 替换为负无穷值（-inf），这样在计算 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)Top logits: tensor([6.7500, 6.2800, 4.5100])
Top positions: tensor([3, 7, 0])随后，我们应用 PyTorch 的 where 函数，将低于 top-3 选择中最低 logit 值的 token 的 logit 值设为负无穷（-inf）。
new_logits = torch.where(
condition=next_token_logits < top_logits[-1], #A
input=torch.tensor(float('-inf')), #B
other=next_token_logits #C
)
print(new_logits)
#A Identifies logits less than the minimum in the top 3
#B Assigns -inf to these lower logits
#C Retains the original logits for all other tokens#A识别小于前 \(3\) 个值中最小值的 logits#B将-inf赋给这些较低的 logits#C为所有其他 token 保留原始 logits
对于由 \(9\) 个 token 构成的词表，下一个 token 的结果 logits 如下：
tensor([4.5100, -inf, -inf, 6.7500, -inf, -inf, -inf, 6.2800, -inf])最后，让我们应用 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])现在，我们可以应用上一节介绍的温度缩放和用于概率采样的 multinomial 函数，从这 \(3\) 个非零概率分数中选择下一个 token，从而生成下一个 token。我们将在下一节通过修改文本生成函数来实现这一点。
5.3.3 修改文本生成函数
前两个小节介绍了两个用于增加 LLM 生成文本多样性的概念：温度采样和 top-k 采样。在本节中，我们会把这些概念结合并加入到前面通过 LLM 生成文本时使用的 generate_simple 函数中，创建一个新的 generate 函数：
def generate(model, idx, max_new_tokens, context_size,
temperature=1.0, top_k=None, eos_id=None):
for _ in range(max_new_tokens): #A
idx_cond = idx[:, -context_size:]
with torch.no_grad():
logits = model(idx_cond)
logits = logits[:, -1, :]
if top_k is not None: #B
top_logits, _ = torch.topk(logits, top_k)
min_val = top_logits[:, -1]
logits = torch.where(
logits < min_val,
torch.tensor(float('-inf')).to(logits.device),
logits
)
if temperature > 0.0: #C
logits = logits / temperature
probs = torch.softmax(logits, dim=-1)
idx_next = torch.multinomial(probs, num_samples=1)
else: #D
idx_next = torch.argmax(logits, dim=-1, keepdim=True)
if idx_next == eos_id: #E
break
idx = torch.cat((idx, idx_next), dim=1)
return idx
#A For-loop is the same as before: Get logits, and only focus on last time step
#B In this new section, we filter logits with top_k sampling
#C This is the new section where we apply temperature scaling
#D Carry out greedy next-token selection as before when temperature scaling is disabled
#E Stop generating early if end-of-sequence token is encountered and eos_id is specified#Afor循环与之前相同：获取 logits，并且只关注最后一个时间步#B在这个新部分中，我们使用 top-k 采样过滤 logits#C这是应用温度缩放的新部分#D当温度缩放被禁用时，像之前一样执行贪婪的下一个 token 选择#E如果遇到序列结束 token 且指定了eos_id，则提前停止生成
现在让我们看看这个新的 generate 函数实际运行起来是什么样：
torch.manual_seed(123)
token_ids = generate(
model=model,
idx=text_to_token_ids("Every effort moves you", tokenizer).to(device),
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 stand to work on surprise, a one of us had gone with random-可以看到，生成的文本与我们在第 5.3 节开头通过 generate_simple 函数生成的文本（"Every effort moves you know," was one of the axioms he laid...!"）非常不同；后者是训练集中的一段记忆文本。
到目前为止，我们已经介绍了如何预训练 LLM，以及如何使用它们生成文本。本章最后两节将讨论如何保存和加载训练好的 LLM，以及如何从 OpenAI 加载预训练权重。
5.4 在 PyTorch 中加载和保存模型权重
在本章中，我们已经讨论了如何用数值方式评估训练进度，并从头开始预训练一个 LLM。尽管这里的 LLM 和数据集都相对较小，这个练习仍然表明，预训练 LLM 的计算成本很高。因此，能够保存 LLM 很重要，这样每次想在新的会话中使用它时，就不需要重新运行训练。
如图 5.16 的本章概览所示，本节会介绍如何保存和加载预训练模型。随后，在下一节中，我们会把来自 OpenAI 的能力更强的预训练 GPT 模型加载到我们的 GPTModel 实例中。
幸运的是，保存 PyTorch 模型相对直接。推荐做法是保存模型所谓的 state_dict，这是一个把每一层映射到其参数的字典，使用 torch.save 函数即可，如下所示：
torch.save(model.state_dict(), "model.pth")在前面的代码中，"model.pth" 是保存 state_dict 的文件名。.pth 扩展名是 PyTorch 文件的惯例，尽管从技术上讲我们可以使用任何文件扩展名。
然后，在通过 state_dict 保存模型权重后，我们可以按如下方式把模型权重加载到一个新的 GPTModel 模型实例中：
model = GPTModel(GPT_CONFIG_124M)
model.load_state_dict(torch.load("model.pth"))
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")
model = GPTModel(GPT_CONFIG_124M)
model.load_state_dict(checkpoint["model_state_dict"])
optimizer = torch.optim.AdamW(model.parameters(), lr=5e-4, weight_decay=0.1)
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
model.train();5.5 加载来自 OpenAI 的预训练权重
此前，出于教学目的，我们使用一个由短篇小说组成的有限数据集，训练了一个小型 GPT-2 模型。这种做法使我们能够专注于基本原理，而不需要大量时间和计算资源。
幸运的是，OpenAI 公开分享了他们的 GPT-2 模型权重，因此我们不需要自己投入数万到数十万美元，在大型语料库上重新训练该模型。
在本节剩余部分中，我们会把这些权重加载到我们的 GPTModel 类中，并使用该模型生成文本。这里的权重指的是权重参数，例如存储在 PyTorch 的 Linear 层和 Embedding 层的 .weight 属性中的参数。在训练模型时，我们前面已经通过 model.parameters() 访问过这些参数。
在接下来的章节中，我们会复用这些预训练权重，把模型微调用于文本分类任务，并让它遵循类似 ChatGPT 的指令。
请注意，OpenAI 最初是通过 TensorFlow 保存 GPT-2 权重的，因此我们必须安装 TensorFlow，才能在 Python 中加载这些权重。此外，下面的代码会使用一个名为 tqdm 的进度条工具来跟踪下载过程，我们也必须安装它。
你可以在终端中执行以下命令来安装这些库：
pip install tensorflow>=2.15.0 tqdm>=4.66下载代码相对较长，大多是样板代码，也并不特别有趣。因此，与其在本章中占用宝贵篇幅讨论从互联网获取文件的 Python 代码，我们直接从本章的在线代码仓库下载 gpt_download.py Python 模块：
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, 63.9kiB/s]
encoder.json: 100%|█████████████████████████| 1.04M/1.04M [00:00<00:00, 2.20MiB/s]
hprams.json: 100%|██████████████████████████| 90.0/90.0 [00:00<00:00, 78.3kiB/s]
model.ckpt.data-00000-of-00001: 100%|███████| 498M/498M [01:09<00:00, 7.16MiB/s]
model.ckpt.index: 100%|█████████████████████| 5.21k/5.21k [00:00<00:00, 3.24MiB/s]
model.ckpt.meta: 100%|██████████████████████| 471k/471k [00:00<00:00, 2.46MiB/s]
vocab.bpe: 100%|████████████████████████████| 456k/456k [00:00<00:00, 1.70MiB/s]前面的代码执行完成后，我们来检查 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)token 嵌入层的权重如下：
[[-0.11010301 ... -0.1363697 0.01506208 0.04531523]
[ 0.04034033 ... 0.08605453 0.00253983 0.04318958]
[-0.12746179 ... 0.08991534 -0.12972379 -0.08785918]
...
[-0.04453601 ... 0.10435229 0.09783269 -0.06952604]
[ 0.1860082 ... -0.09625227 0.07847701 -0.02245961]
[ 0.05135201 ... 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 模型的整体架构保持相同，只是不同架构元素重复的次数不同，嵌入大小也不同。本章剩余代码同样兼容这些更大的模型。
把 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：
model_name = "gpt2-small (124M)"
NEW_CONFIG = GPT_CONFIG_124M.copy()
NEW_CONFIG.update(model_configs[model_name])细心的读者可能还记得，我们之前使用的是 \(256\)-token 长度，但 OpenAI 的原始 GPT-2 模型是用 \(1,024\)-token 长度训练的，因此我们必须相应地更新 NEW_CONFIG：
NEW_CONFIG.update({"context_length": 1024})此外，OpenAI 在多头注意力模块的线性层中使用了偏置向量，用来实现查询、键和值矩阵的计算。LLM 现在通常不再使用偏置向量，因为它们不会提升建模性能，因此并非必要。不过，由于我们正在使用预训练权重，为保持一致，我们需要匹配这些设置并启用这些偏置向量：
NEW_CONFIG.update({"qkv_bias": True})现在，我们可以使用更新后的 NEW_CONFIG 字典初始化一个新的 GPTModel 实例：
gpt = GPTModel(NEW_CONFIG)
gpt.eval()默认情况下，GPTModel 实例会使用用于预训练的随机权重初始化。使用 OpenAI 模型权重的最后一步，是用我们加载到 params 字典中的权重覆盖这些随机权重。
为此，我们首先定义一个小型 assign 工具函数，它会检查两个张量或数组（左侧和右侧）是否具有相同维度或形状，并把右侧张量作为可训练的 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 中：
import numpy as np
def load_weights_into_gpt(gpt, params):
gpt.pos_emb.weight = assign(gpt.pos_emb.weight, params['wpe']) #A
gpt.tok_emb.weight = assign(gpt.tok_emb.weight, params['wte'])
for b in range(len(params["blocks"])): #B
q_w, k_w, v_w = np.split( #C
(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"]) #D#A将模型的位置嵌入权重和 token 嵌入权重设置为params中指定的权重。#B遍历模型中的每个 transformer 块。#C使用np.split函数把注意力权重和偏置权重等分为三部分，分别用于查询、键和值组件。#DOpenAI 的原始 GPT-2 模型在输出层复用了 token 嵌入权重，以减少总参数数量；这个概念称为权重绑定。
在 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),
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 new way to practice something!
What makes us want to be on top of that?我们可以确信模型权重加载正确，因为模型能够生成连贯文本。这个过程中的一个很小错误都会导致模型失败。
在接下来的章节中，我们会继续使用这个预训练模型，并对它进行微调，使其能够分类文本和遵循指令。
5.6 小结
- 当 LLM 生成文本时，它们一次输出一个 token。
- 默认情况下，生成下一个 token 的方式是把模型输出转换为概率分数，并从词表中选择与最高概率分数对应的 token；这称为“贪婪解码”。
- 使用概率采样和温度缩放，我们可以影响生成文本的多样性和连贯性。
- 训练集和验证集损失可以用于衡量 LLM 在训练期间生成文本的质量。
- 预训练 LLM 涉及改变其权重，以最小化训练损失。
- LLM 的训练循环本身是深度学习中的标准流程，使用常规的交叉熵损失和 AdamW 优化器。
- 在大型文本语料库上预训练 LLM 非常耗时且消耗资源，因此我们可以加载 OpenAI 公开提供的权重，作为自己在大型数据集上预训练模型的替代方案。
章节范围：原 PDF 物理第 155-203 页。图片从同一页段抽取并嵌入本文。