第 5 章 在无标注数据上预训练
第 5 章 在无标注数据上预训练
本章来源：本文翻译整理自 LLMs-from-scratch 仓库的 ch05/01_main-chapter-code/ch05.ipynb，原书为 Sebastian Raschka《Build a Large Language Model (From Scratch)》。
本章要做什么
在本章，我们实现训练循环和基本模型评估的代码来预训练一个 LLM。在本章末尾，我们还会把 OpenAI 公开可用的预训练权重加载进我们的模型。具体包括：评估生成式文本模型、计算文本生成损失（交叉熵和困惑度）、计算训练和验证集损失、训练 LLM、控制随机性的解码策略（温度缩放和 top-k 采样）、在 PyTorch 里保存和加载模型权重、从 OpenAI 加载预训练权重。
本 notebook 使用的包
from importlib.metadata import version
pkgs = ["matplotlib",
"numpy",
"tiktoken",
"torch",
"tensorflow" # For OpenAI's pretrained weights
]
for p in pkgs:
print(f"{p} version: {version(p)}")5.1 评估生成式文本模型
我们以简要回顾用上一章的代码初始化 GPT 模型来开始本节。然后，我们讨论 LLM 的基本评估指标。最后，在本节我们把这些评估指标应用到一个训练集和一个验证集上。
5.1.1 用 GPT 生成文本
我们用上一章的代码初始化一个 GPT 模型：
import torch
from previous_chapters import GPTModel
# If the `previous_chapters.py` file is not available locally,
# you can import it from the `llms-from-scratch` PyPI package.
# For details, see: https://github.com/rasbt/LLMs-from-scratch/tree/main/pkg
# E.g.,
# from llms_from_scratch.ch04 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)
model = GPTModel(GPT_CONFIG_124M)
model.eval(); # Disable dropout during inference上面我们用 0.1 的 dropout，但如今训练 LLM 不用 dropout 也相对常见。现代 LLM 在查询、键和值矩阵的 nn.Linear 层里也不用偏置向量（不像早期的 GPT 模型），这是通过设置 "qkv_bias": False 实现的。
我们把上下文长度（context_length）减少到只有 256 个 token，以减少训练模型所需的计算资源，而原始的 1.24 亿参数 GPT-2 模型使用 1024 个 token。
- 这是为了让更多读者能跟着在笔记本电脑上执行代码示例。
- 不过，请随意把
context_length增加到 1024 个 token（这不需要任何代码改动）。 - 我们稍后也会从预训练权重加载一个
context_length为 1024 的模型。
接下来，我们用上一章的 generate_text_simple 函数生成文本。此外，我们定义两个便捷函数 text_to_token_ids 和 token_ids_to_text，用于在 token 和文本表示之间转换，本章通篇都会用到：
import tiktoken
from previous_chapters import generate_text_simple
# Alternatively:
# from llms_from_scratch.ch04 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))正如我们上面看到的，模型没有产生好的文本，因为它还没有被训练。我们怎么以数字形式度量或捕获"好的文本"是什么，以便在训练期间跟踪它？下一小节介绍计算生成输出损失指标的度量方法，我们可以用它来衡量训练进度。接下来关于微调 LLM 的章节还会介绍其它衡量模型质量的方法。
5.1.2 计算文本生成损失：交叉熵和困惑度
假设我们有一个 inputs 张量，包含 2 个训练示例（行）的 token ID。对应 inputs，targets 包含我们希望模型生成的期望 token ID。注意 targets 是 inputs 向右移一位，正如第 2 章实现数据加载器时解释的。
inputs = torch.tensor([[16833, 3626, 6100], # ["every effort moves",
[40, 1107, 588]]) # "I really like"]
targets = torch.tensor([[3626, 6100, 345 ], # [" effort moves you",
[1107, 588, 11311]]) # " really like chocolate"]把 inputs 喂给模型，我们得到 2 个输入示例（各含 3 个 token）的 logits 向量。每个 token 是一个 50,257 维向量，对应词汇表的大小。应用 softmax 函数，我们可以把 logits 张量转成同维度的概率分数张量：
with torch.no_grad():
logits = model(inputs)
probas = torch.softmax(logits, dim=-1) # Probability of each token in vocabulary
print(probas.shape) # Shape: (batch_size, num_tokens, vocab_size)下面的图用很小的词汇表做说明，概述了我们如何把概率分数转回文本，这是我们在上一章末尾讨论过的。
正如上一章讨论的，我们可以应用 argmax 函数把概率分数转成预测的 token ID。上面的 softmax 函数为每个 token 产生一个 50,257 维向量；argmax 函数返回这个向量里最高概率分数的位置，也就是给定 token 的预测 token ID。
因为我们有 2 个输入批次，各含 3 个 token，我们得到 2 乘 3 的预测 token ID：
token_ids = torch.argmax(probas, dim=-1, keepdim=True)
print("Token IDs:\n", token_ids)如果我们解码这些 token，我们会发现它们和我们希望模型预测的 token，也就是目标 token，相当不同：
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 概率如下：
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)我们想最大化所有这些值，让它们接近概率 1。在数学优化里，最大化概率分数的对数比最大化概率分数本身更容易；这超出本书范围，但作者录了一节课讲更多细节：L8.2 Logistic Regression Loss Function。
# Compute logarithm of all token probabilities
log_probas = torch.log(torch.cat((target_probas_1, target_probas_2)))
print(log_probas)接下来，我们计算平均对数概率：
# Calculate the average probability for each token
avg_log_probas = torch.mean(log_probas)
print(avg_log_probas)目标是通过优化模型权重，让这个平均对数概率尽可能大。由于对数的原因，最大的可能值是 0，而我们目前离 0 很远。
在深度学习中，与最大化平均对数概率相反，标准惯例是最小化负的平均对数概率值；在我们的例子里，深度学习里我们不会最大化 -10.7722 让它接近 0，而是最小化 10.7722 让它接近 0。-10.7722 的负值，也就是 10.7722，在深度学习里也叫交叉熵损失（cross-entropy loss）。
neg_avg_log_probas = avg_log_probas * -1
print(neg_avg_log_probas)PyTorch 已经实现了一个 cross_entropy 函数，它执行前面的步骤。
在我们应用 cross_entropy 函数之前，让我们检查一下 logits 和目标的形状。
# 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)对 PyTorch 里的 cross_entropy 函数，我们想通过结合 batch 维度来展平这些张量：
logits_flat = logits.flatten(0, 1)
targets_flat = targets.flatten()
print("Flattened logits:", logits_flat.shape)
print("Flattened targets:", targets_flat.shape)注意目标就是 token ID，它们也表示 logits 张量里我们想要最大化的索引位置。PyTorch 里的 cross_entropy 函数会自动在内部对 logits 里那些要被最大化的 token 索引应用 softmax 和对数概率计算。
loss = torch.nn.functional.cross_entropy(logits_flat, targets_flat)
print(loss)一个与交叉熵损失相关的概念是 LLM 的困惑度（perplexity）。困惑度就是交叉熵损失的指数。
perplexity = torch.exp(loss)
print(perplexity)困惑度通常被认为更容易解释，因为它可以理解为模型每一步不确定的有效词汇表大小（在上面的例子里，那是 48,725 个词或 token）。换句话说，困惑度提供了一个度量，衡量模型预测的概率分布和数据集里词的实际分布匹配得多好。和损失类似，更低的困惑度表示模型预测更接近实际分布。
5.1.3 计算训练和验证集损失
我们用相对较小的数据集训练 LLM（实际上只是一篇短篇小说）。原因是：
- 你可以在没有合适 GPU 的笔记本电脑上用几分钟运行代码示例。
- 训练相对较快地完成（分钟而不是周），这对教育目的很好。
- 我们用公有领域的文本，它可以在不违反任何使用权利或膨胀仓库大小的情况下包含在这个 GitHub 仓库里。
比如，Llama 2 7B 需要 184,320 个 A100 GPU 小时在 2 万亿 token 上训练。在写这篇文章时，AWS 上 8xA100 云服务器的小时成本大约是 \$30。所以，通过粗略计算，训练这个 LLM 会花 184,320 / 8 * \$30 = \$690,000。
下面，我们用第 2 章用过的同一个数据集：
import os
import requests
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):
response = requests.get(url, timeout=30)
response.raise_for_status()
text_data = response.text
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()一个快速检查，通过打印前 99 个字符确认文本加载正常：
# First 99 characters
print(text_data[:99])# Last 99 characters
print(text_data[-99:])total_characters = len(text_data)
total_tokens = len(tokenizer.encode(text_data))
print("Characters:", total_characters)
print("Tokens:", total_tokens)有 5,145 个 token，这篇文本对训练 LLM 来说非常短，但再说一次，这是为了教育目的（我们稍后也会加载预训练权重）。
接下来，我们把数据集分成训练集和验证集，用第 2 章的数据加载器准备 LLM 训练的批次。为了可视化，下面的图假设 max_length=6，但对训练加载器，我们把 max_length 设为等于 LLM 支持的上下文长度。下图为了简单只显示输入 token。因为我们训练 LLM 预测文本里的下一个词，目标看起来和这些输入一样，只是目标向右移一位。
from previous_chapters import create_dataloader_v1
# Alternatively:
# from llms_from_scratch.ch02 import create_dataloader_v1
# 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
)# Sanity check
if total_tokens * (train_ratio) < GPT_CONFIG_124M["context_length"]:
print("Not enough tokens for the training loader. "
"Try to lower the `GPT_CONFIG_124M['context_length']` or "
"increase the `training_ratio`")
if total_tokens * (1-train_ratio) < GPT_CONFIG_124M["context_length"]:
print("Not enough tokens for the validation loader. "
"Try to lower the `GPT_CONFIG_124M['context_length']` or "
"decrease the `training_ratio`")我们用相对较小的批量大小来减少计算资源需求，也因为数据集本身非常小。比如，Llama 2 7B 是用批量大小 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)另一个可选的检查，确认 token 大小在预期的范围内：
train_tokens = 0
for input_batch, target_batch in train_loader:
train_tokens += input_batch.numel()
val_tokens = 0
for input_batch, target_batch in val_loader:
val_tokens += input_batch.numel()
print("Training tokens:", train_tokens)
print("Validation tokens:", val_tokens)
print("All tokens:", train_tokens + val_tokens)接下来，我们实现一个工具函数来计算给定批次的交叉熵损失。此外，我们实现第二个工具函数来计算数据加载器里用户指定批次数的损失。
def calc_loss_batch(input_batch, target_batch, model, device):
input_batch, target_batch = input_batch.to(device), target_batch.to(device)
logits = model(input_batch)
loss = torch.nn.functional.cross_entropy(logits.flatten(0, 1), target_batch.flatten())
return loss
def calc_loss_loader(data_loader, model, device, num_batches=None):
total_loss = 0.
if len(data_loader) == 0:
return float("nan")
elif num_batches is None:
num_batches = len(data_loader)
else:
# Reduce the number of batches to match the total number of batches in the data loader
# if num_batches exceeds the number of batches in the data loader
num_batches = min(num_batches, len(data_loader))
for i, (input_batch, target_batch) in enumerate(data_loader):
if i < num_batches:
loss = calc_loss_batch(input_batch, target_batch, model, device)
total_loss += loss.item()
else:
break
return total_loss / num_batches如果你的机器有支持 CUDA 的 GPU，LLM 会在 GPU 上训练，不需要改任何代码。通过 device 设置，我们确保数据被加载到和 LLM 模型相同的设备上。
if torch.cuda.is_available():
device = torch.device("cuda")
elif torch.backends.mps.is_available():
# Use PyTorch 2.9 or newer for stable mps results
major, minor = map(int, torch.__version__.split(".")[:2])
if (major, minor) >= (2, 9):
device = torch.device("mps")
else:
device = torch.device("cpu")
else:
device = torch.device("cpu")
print(f"Using {device} device.")
model.to(device) # no assignment model = model.to(device) necessary for nn.Module classes
torch.manual_seed(123) # For reproducibility due to the shuffling in the data loader
with torch.no_grad(): # Disable gradient tracking for efficiency because we are not training, yet
train_loss = calc_loss_loader(train_loader, model, device)
val_loss = calc_loss_loader(val_loader, model, device)
print("Training loss:", train_loss)
print("Validation loss:", val_loss)5.2 训练一个 LLM
在本节，我们终于实现训练 LLM 的代码。我们聚焦一个简单的训练函数（如果你有兴趣用更高级的技术增强这个训练函数，比如学习率预热、余弦退火和梯度裁剪，请参见附录 D）。
def train_model_simple(model, train_loader, val_loader, optimizer, device, num_epochs,
eval_freq, eval_iter, start_context, tokenizer):
# Initialize lists to track losses and tokens seen
train_losses, val_losses, track_tokens_seen = [], [], []
tokens_seen, global_step = 0, -1
# Main training loop
for epoch in range(num_epochs):
model.train() # Set model to training mode
for input_batch, target_batch in train_loader:
optimizer.zero_grad() # Reset loss gradients from previous batch iteration
loss = calc_loss_batch(input_batch, target_batch, model, device)
loss.backward() # Calculate loss gradients
optimizer.step() # Update model weights using loss gradients
tokens_seen += input_batch.numel()
global_step += 1
# Optional evaluation step
if global_step % eval_freq == 0:
train_loss, val_loss = evaluate_model(
model, train_loader, val_loader, device, eval_iter)
train_losses.append(train_loss)
val_losses.append(val_loss)
track_tokens_seen.append(tokens_seen)
print(f"Ep {epoch+1} (Step {global_step:06d}): "
f"Train loss {train_loss:.3f}, Val loss {val_loss:.3f}")
# Print a sample text after each epoch
generate_and_print_sample(
model, tokenizer, device, start_context
)
return train_losses, val_losses, track_tokens_seen
def evaluate_model(model, train_loader, val_loader, device, eval_iter):
model.eval()
with torch.no_grad():
train_loss = calc_loss_loader(train_loader, model, device, num_batches=eval_iter)
val_loss = calc_loss_loader(val_loader, model, device, num_batches=eval_iter)
model.train()
return train_loss, val_loss
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()现在，让我们用上面定义的训练函数训练 LLM：
# 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.")注意，你的电脑上可能得到略有不同的损失值，如果大致相似（训练损失低于 1、验证损失低于 7），这不是需要担心的问题。小的差异通常可能是由于不同的 GPU 硬件和 CUDA 版本，或更新的 PyTorch 版本里的小改动。即使你在 CPU 上运行示例，也可能观察到细微差异；一个可能的原因是 nn.Dropout 在不同操作系统上、取决于 PyTorch 的编译方式而有不同行为，如这里 PyTorch issue tracker 所讨论的。
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)看上面的结果，我们可以看到模型一开始生成难以理解的词串，而到接近结尾时，它能产生语法大致正确的句子。然而，基于训练和验证集损失，我们可以看到模型开始过拟合。如果我们检查它接近结尾写的一些段落，我们会发现它们逐字出现在训练集里，它只是记住了训练数据。稍后我们会覆盖能在一定程度上缓解这种记忆的解码策略。
注意，这里的过拟合是因为我们的训练集非常非常小，而且我们迭代它太多次了。这里的 LLM 训练主要用于教育目的；我们主要想看到模型能学会产生连贯的文本。与其花几周或几个月在大量昂贵硬件上训练这个模型，我们稍后加载预训练权重。
如果你有兴趣用更高级的技术增强这个训练函数，比如学习率预热、余弦退火和梯度裁剪，请参见附录 D。
如果你有兴趣用更大的训练数据集和更长的训练运行，参见 ../03_bonus_pretraining_on_gutenberg。
5.3 控制随机性的解码策略
对像我们上面训练的相对小的 LLM 来说，推理相对便宜，所以如果你上面用 GPU 训练它，推理没有用 GPU 的必要。用我们之前在简单训练函数里用过的 generate_text_simple 函数（来自上一章），我们可以一次一个词（或 token）地生成新文本。如 5.1.2 节解释的，下一个生成的 token 就是词汇表所有 token 里对应最大概率分数的那个 token。
# NEW: use CPU here as inference is cheap with
# this model and to ensure readers get same results in the
# remaining sections of this book
inference_device = torch.device("cpu")
model.to(inference_device)
model.eval()
tokenizer = tiktoken.get_encoding("gpt2")
token_ids = generate_text_simple(
model=model,
idx=text_to_token_ids("Every effort moves you", tokenizer).to(inference_device),
max_new_tokens=25,
context_size=GPT_CONFIG_124M["context_length"]
)
print("Output text:\n", token_ids_to_text(token_ids, tokenizer))即使我们多次执行上面的 generate_text_simple 函数，LLM 也总是生成相同的输出。我们现在引入两个概念，也就是所谓的解码策略，来修改 generate_text_simple：温度缩放（temperature scaling）和 top-k 采样。这些将允许模型控制生成文本的随机性和多样性。
5.3.1 温度缩放
之前，我们总是用 torch.argmax 采样概率最高的 token 作为下一个 token。为了增加多样性，我们可以用 torch.multinomial(probs, num_samples=1) 从概率分布里采样下一个 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()}
# Suppose input is "every effort moves you", and the LLM
# returns the following logits for the next token:
next_token_logits = torch.tensor(
[4.51, 0.89, -1.90, 6.75, 1.63, -1.62, -1.89, 6.28, 1.79]
)
probas = torch.softmax(next_token_logits, dim=0)
next_token_id = torch.argmax(probas).item()
# The next generated token is then as follows:
print(inverse_vocab[next_token_id])torch.manual_seed(123)
next_token_id = torch.multinomial(probas, num_samples=1).item()
print(inverse_vocab[next_token_id])与用 torch.argmax 确定最可能的 token 不同，我们用 torch.multinomial(probas, num_samples=1) 通过从 softmax 分布采样来确定最可能的 token。为了说明，让我们看看用原始 softmax 概率采样下一个 token 1,000 次会发生什么：
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)我们可以通过一个叫温度缩放（temperature scaling）的概念控制分布和选择过程。"温度缩放"只是把 logits 除以一个大于 0 的数的花哨说法。大于 1 的温度在应用 softmax 后会产生更均匀分布的 token 概率。小于 1 的温度在应用 softmax 后会产生更自信（更尖锐或更多峰）的分布。
def softmax_with_temperature(logits, temperature):
scaled_logits = logits / temperature
return torch.softmax(scaled_logits, dim=0)
# Temperature values
temperatures = [1, 0.1, 5] # Original, higher confidence, and lower confidence
# Calculate scaled probabilities
scaled_probas = [softmax_with_temperature(next_token_logits, T) for T in temperatures]# Plotting
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.savefig("temperature-plot.pdf")
plt.show()我们可以看到，通过温度 0.1 的重新缩放产生更尖锐的分布，接近 torch.argmax，因此最可能的词几乎总是被选中：
print_sampled_tokens(scaled_probas[1])通过温度 5 重新缩放的概率分布更均匀：
print_sampled_tokens(scaled_probas[2])假设 LLM 输入 "every effort moves you"，用上面的方法有时会产生无意义的文本，比如 "every effort moves you pizza"，3.2% 的时间（1000 次里 32 次）。
5.3.2 Top-k 采样
为了能用更高温度来增加输出多样性、并减少无意义句子的概率，我们可以把采样的 token 限制到最可能的 top-k 个 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)new_logits = torch.where(
condition=next_token_logits < top_logits[-1],
input=torch.tensor(float("-inf")),
other=next_token_logits
)
print(new_logits)注意：
前一个代码单元的一个替代、略更高效的实现如下：
new_logits = torch.full_like( # create tensor containing -inf values next_token_logits, -torch.inf ) new_logits[top_pos] = next_token_logits[top_pos] # copy top k values into the -inf tensor
更多细节，见 https://github.com/rasbt/LLMs-from-scratch/discussions/326
topk_probas = torch.softmax(new_logits, dim=0)
print(topk_probas)5.3.3 修改文本生成函数
前面两个小节介绍了温度采样和 top-k 采样。让我们用这两个概念修改第 4 章的 generate_text_simple 函数，创建一个新的 generate 函数：
def generate(model, idx, max_new_tokens, context_size, temperature=0.0, top_k=None, eos_id=None):
# For-loop is the same as before: Get logits, and only focus on last time step
for _ in range(max_new_tokens):
idx_cond = idx[:, -context_size:]
with torch.no_grad():
logits = model(idx_cond)
logits = logits[:, -1, :]
# New: Filter logits with top_k sampling
if top_k is not None:
# Keep only top_k values
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)
# New: Apply temperature scaling
if temperature > 0.0:
logits = logits / temperature
# New (not in book): numerical stability tip to get equivalent results on mps device
# subtract rowwise max before softmax
logits = logits - logits.max(dim=-1, keepdim=True).values
# Apply softmax to get probabilities
probs = torch.softmax(logits, dim=-1) # (batch_size, context_len)
# Sample from the distribution
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:
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
break
# Same as before: append sampled index to the running sequence
idx = torch.cat((idx, idx_next), dim=1) # (batch_size, num_tokens+1)
return idxtorch.manual_seed(123)
token_ids = generate(
model=model,
idx=text_to_token_ids("Every effort moves you", tokenizer).to(inference_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))5.4 在 PyTorch 里保存和加载模型权重
训练 LLM 计算上很昂贵，所以能保存和加载 LLM 权重很关键。
PyTorch 里推荐的方式是通过把 torch.save 函数应用到 .state_dict() 方法来保存模型权重，也就是所谓的 state_dict：
torch.save(model.state_dict(), "model.pth")然后我们可以把模型权重加载进一个新的 GPTModel 模型实例，如下所示：
model = GPTModel(GPT_CONFIG_124M)
if torch.cuda.is_available():
device = torch.device("cuda")
elif torch.backends.mps.is_available():
# Use PyTorch 2.9 or newer for stable mps results
major, minor = map(int, torch.__version__.split(".")[:2])
if (major, minor) >= (2, 9):
device = torch.device("mps")
else:
device = torch.device("cpu")
print("Device:", device)
model.load_state_dict(torch.load("model.pth", map_location=device, weights_only=True))
model.eval();用 Adam 或 AdamW 这类自适应优化器而不是常规 SGD 训练 LLM 很常见。这些自适应优化器为每个模型权重存储额外参数，所以如果我们计划稍后继续预训练，保存它们也是合理的：
torch.save({
"model_state_dict": model.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
},
"model_and_optimizer.pth"
)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.5 从 OpenAI 加载预训练权重
之前，我们只用一篇非常短的短篇小说书训练了一个小型 GPT-2 模型，用于教育目的。感兴趣的读者还可以在 ../03_bonus_pretraining_on_gutenberg 找到在完整 Project Gutenberg 图书语料上的更长预训练运行。
幸运的是，我们不必花几万到几十万美元在大型预训练语料上预训练模型，而是可以加载 OpenAI 提供的预训练权重。
⚠️ 注意：有些用户可能在本节遇到 TensorFlow 兼容性问题，尤其是在某些 Windows 系统上。这里需要 TensorFlow 只是为了加载原始的 OpenAI GPT-2 权重文件，然后我们把它转换成 PyTorch。如果你遇到和 TensorFlow 相关的问题，你可以用下面的替代代码，而不是本节剩余部分的代码。这个替代方案基于预先转换的 PyTorch 权重，是用前一节描述的相同转换过程创建的。详情参见 notebook：../02_alternative_weight_loading/weight-loading-pytorch.ipynb。
file_name = "gpt2-small-124M.pth"
# file_name = "gpt2-medium-355M.pth"
# file_name = "gpt2-large-774M.pth"
# file_name = "gpt2-xl-1558M.pth"
url = f"https://huggingface.co/rasbt/gpt2-from-scratch-pytorch/resolve/main/{file_name}"
if not os.path.exists(file_name):
urllib.request.urlretrieve(url, file_name)
print(f"Downloaded to {file_name}")
gpt = GPTModel(BASE_CONFIG)
gpt.load_state_dict(torch.load(file_name, weights_only=True))
gpt.eval()
if torch.cuda.is_available():
device = torch.device("cuda")
elif torch.backends.mps.is_available():
# Use PyTorch 2.9 or newer for stable mps results
major, minor = map(int, torch.__version__.split(".")[:2])
if (major, minor) >= (2, 9):
device = torch.device("mps")
else:
device = torch.device("cpu")
gpt.to(device);
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))首先，一些样板代码从 OpenAI 下载文件并把权重加载进 Python。因为 OpenAI 用了 TensorFlow，我们不得不安装并使用 TensorFlow 来加载权重；tqdm 是一个进度条库。取消注释并运行下一个单元来安装所需库。
# pip install tensorflow tqdmprint("TensorFlow version:", version("tensorflow"))
print("tqdm version:", version("tqdm"))# Relative import from the gpt_download.py contained in this folder
from gpt_download import download_and_load_gpt2
# Alternatively:
# from llms_from_scratch.ch05 import download_and_load_gpt2注意
- 在非常罕见的情况下，上面的代码单元可能导致
zsh: illegal hardware instruction python错误，这可能是因为你机器上的 TensorFlow 安装问题。 - 一位读者发现在这个特定情况下通过
conda安装 TensorFlow 解决了问题，如这里提到的。 - 你可以在这个补充的 Python 安装教程 里找到更多说明。
然后我们可以如下下载 1.24 亿参数模型的权重：
settings, params = download_and_load_gpt2(model_size="124M", models_dir="gpt2")print("Settings:", settings)print("Parameter dictionary keys:", params.keys())print(params["wte"])
print("Token embedding weight tensor dimensions:", params["wte"].shape)另外，"355M"、"774M" 和 "1558M" 也是受支持的 model_size 参数。这些不同大小模型之间的区别总结在下图里：
上面，我们把 124M GPT-2 模型的权重加载进 Python，但我们还需要把它们转进我们的 GPTModel 实例。首先，我们初始化一个新的 GPTModel 实例。
注意，原始 GPT 模型初始化多头注意力模块里查询、键和值矩阵的线性层时带有偏置向量，这不是必需也不推荐；然而，为了能正确加载权重，我们在实现里也不得不通过设置 qkv_bias 为 True 来启用它们。我们也使用原始 GPT-2 模型使用的 1024 token 上下文长度。
# 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();下一个任务是把 OpenAI 权重赋给我们 GPTModel 实例里对应的权重张量。
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))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);如果模型加载正确，我们可以用它和之前的 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))我们知道我们正确加载了模型权重，因为模型能生成连贯的文本；如果我们犯了哪怕一个小错误，模型都无法做到这一点。
- 从 Hugging Face Hub 加载权重的另一种方式，参见 ../02_alternative_weight_loading。
- 如果你有兴趣看 GPT 架构和 Llama 架构（Meta AI 开发的热门 LLM）怎么比较，参见 ../07_gpt_to_llama 的补充内容。
总结与要点
- 参见 ./gpt_train.py 脚本，一个用于训练的自包含脚本。
- ./gpt_generate.py 脚本从 OpenAI 加载预训练权重，基于提示生成文本。
- 你可以在 ./exercise-solutions.ipynb 找到练习解答。
关键概念
- 交叉熵损失（cross-entropy loss）：负平均对数概率，衡量模型预测分布和实际分布的距离。
- 困惑度（perplexity）：交叉熵损失的指数，可理解为模型每步不确定的有效词汇表大小。
- 温度缩放（temperature scaling）：把 logits 除以温度；>1 更均匀，<1 更尖锐。
- top-k 采样：把采样限制到最可能的 k 个 token，减少无意义输出。
- 贪婪解码 vs 采样：argmax 总是选最高概率；multinomial 从分布采样。
- state_dict：PyTorch 保存/加载模型权重的方式。
- 预训练权重：直接加载 OpenAI 的 GPT-2 权重，省去昂贵的预训练。
练习
练习解答见仓库的 ch05/01_main-chapter-code/exercise-solutions.ipynb。