从零构建LLM 第5章:预训练
1. 预训练的目标:下一个 Token 预测
LLM 的预训练目标极其简洁:给定前 t 个 token,预测第 t+1 个 token。这就是第2章中滑动窗口生成的 (input, target) 对的用武之地。
衡量预测好坏的指标是交叉熵损失 (Cross-Entropy Loss):模型输出一个概率分布(50257 维向量),真实答案是一个 token ID。交叉熵衡量的是模型的概率分布与真实分布之间的「距离」— 损失越小,说明模型给正确 token 分配的概率越高。
def calc_loss_batch(input_batch, target_batch, model, device):
input_batch = input_batch.to(device)
target_batch = target_batch.to(device)
logits = model(input_batch) # [batch, seq_len, vocab_size]
# 交叉熵损失:模型预测 vs 真实的下一个 token
loss = torch.nn.functional.cross_entropy(
logits.flatten(0, 1), # [batch*seq_len, 50257]
target_batch.flatten() # [batch*seq_len]
)
return loss
def calc_loss_loader(data_loader, model, device, num_batches=None):
"""在整个 DataLoader 上计算平均损失"""
total_loss = 0.
if num_batches is None:
num_batches = len(data_loader)
else:
num_batches = min(num_batches, len(data_loader))
for i, (input_batch, target_batch) in enumerate(data_loader):
if i >= num_batches:
break
loss = calc_loss_batch(input_batch, target_batch, model, device)
total_loss += loss.item()
return total_loss / num_batches
ℹ️ 从损失到困惑度:困惑度 (Perplexity) = eloss。初始随机模型的困惑度约等于词汇表大小 (50257),意味着模型在「随机猜」。训练后的 GPT-2 困惑度可以降到 20 以下,意味着每个位置模型只在约 20 个候选词之间犹豫。
2. 训练循环:让模型学会语言
训练循环是深度学习的核心模式:前向传播计算损失 → 反向传播计算梯度 → 优化器更新权重。书中用 AdamW 优化器(Adam + 权重衰减),这也是几乎所有现代 LLM 的标配。
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 = [], [], []
tokens_seen, global_step = 0, -1
for epoch in range(num_epochs):
model.train()
for input_batch, target_batch in train_loader:
optimizer.zero_grad() # 清空梯度
loss = calc_loss_batch(input_batch, target_batch,
model, device)
loss.backward() # 反向传播
optimizer.step() # 更新权重
tokens_seen += input_batch.numel()
global_step += 1
# 定期评估训练集和验证集损失
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)
print(f"Ep {epoch+1} Step {global_step:06d}: "
f"Train loss {train_loss:.3f}, Val loss {val_loss:.3f}")
# 每个 epoch 结束,生成示例文本观察模型能力变化
generate_and_print_sample(model, tokenizer, device, start_context)
return train_losses, val_losses, track_tokens_seen
书中用一段很短的文本(”The Verdict” 短篇小说)做演示。训练 10 个 epoch 后,模型从输出乱码逐渐进化到能生成语法基本正确的句子 — 虽然内容还很有限(因为训练数据太小),但足以展示预训练的工作原理。
3. 文本生成:温度与 Top-k 采样
模型训练好之后,如何用它生成文本?核心在于如何从模型输出的概率分布中选择下一个 token。书中介绍了三种策略:
| 策略 | temperature | top_k | 特点 |
|---|---|---|---|
| 贪心解码 | 0 | — | 始终选概率最高的 token,确定性但容易重复 |
| 温度采样 | 0.1–2.0 | — | 低温接近贪心,高温更随机多样 |
| Top-k 采样 | 0.7–1.0 | 20–50 | 只从概率最高的 k 个 token 中采样 |
def generate(model, idx, max_new_tokens, context_size,
temperature=0.0, top_k=None, eos_id=None):
for _ in range(max_new_tokens):
idx_cond = idx[:, -context_size:] # 截断到上下文长度
with torch.no_grad():
logits = model(idx_cond)
logits = logits[:, -1, :] # 只取最后位置的 logits
# Top-k 过滤:只保留概率最高的 k 个候选
if top_k is not None:
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
)
# 温度缩放:temperature > 1 → 更平均;< 1 → 更尖锐
if temperature > 0.0:
logits = logits / temperature
probs = torch.softmax(logits, dim=-1)
idx_next = torch.multinomial(probs, num_samples=1)
else:
idx_next = torch.argmax(logits, dim=-1, keepdim=True)
if eos_id is not None and idx_next.item() == eos_id:
break
idx = torch.cat((idx, idx_next), dim=1)
return idx
⚠️ Temperature 的数学直觉:logits 除以 temperature 后再 softmax。当 temperature → 0 时,最大 logit 对应的概率趋近 1(贪心);当 temperature → ∞ 时,所有 token 的概率趋近相等(纯随机)。实际使用中 0.7–1.0 是常见范围。
4. 学习率调度:预热与余弦衰减
学习率 (Learning Rate) 是训练中最关键的超参数。用一个固定的学习率往往不够好 — 书中介绍了现代 LLM 普遍采用的线性预热 + 余弦衰减策略:
- 预热阶段 (Warmup):学习率从 0 线性增长到峰值。避免训练初期步子太大破坏随机初始化的权重分布
- 余弦衰减 (Cosine Decay):学习率按余弦曲线从峰值平滑下降到最小值。后期用小学习率精细调优
import math
total_steps = len(train_loader) * num_epochs
warmup_steps = int(0.1 * total_steps) # 前 10% 步数用于预热
peak_lr = 4e-4
min_lr = peak_lr * 0.1 # 最终衰减到峰值的 10%
def get_lr(step):
# 阶段1: 线性预热 (0 → peak_lr)
if step < warmup_steps:
return peak_lr * (step + 1) / warmup_steps
# 阶段2: 余弦衰减 (peak_lr → min_lr)
progress = (step - warmup_steps) / (total_steps - warmup_steps)
return min_lr + 0.5 * (peak_lr - min_lr) * (
1 + math.cos(math.pi * progress)
)
# 在训练循环中手动设置学习率
optimizer = torch.optim.AdamW(model.parameters(), weight_decay=0.1)
for step in range(total_steps):
lr = get_lr(step)
for param_group in optimizer.param_groups:
param_group["lr"] = lr
# ... 正常训练步骤(forward → backward → step)
书中还演示了如何加载 OpenAI 发布的 GPT-2 预训练权重。需要做的就是将 OpenAI 的权重名称映射到我们自己实现的模型参数上,然后逐一赋值。加载权重后模型就能直接用于文本生成 — 这验证了我们的实现与原版 GPT-2 架构完全一致。
5. 为什么这很重要
本章覆盖了从「随机权重」到「能生成文本」的完整过程,以下是几个值得深入理解的要点:
- 预训练 vs 微调:预训练是一个通用阶段 — 模型在大量文本上学习语言的普遍规律,不针对任何特定任务。后面第6-7章的微调则是让模型「专精」某个具体任务(分类、对话等)。这种「先通用后专精」的范式是现代 AI 的基石。
- Scaling 的前提:训练循环的效率直接决定了模型能否扩大规模。AdamW、学习率调度、梯度累积、混合精度等技巧都是让更大模型能在有限硬件上训练的关键。
- 生成策略的工程价值:理解 temperature 和 top-k 对于使用 LLM API 至关重要。代码生成通常用低温(确定性),创意写作用高温(多样性),对话场景则需要 top-k/top-p 来平衡。
- Loss 曲线诊断:训练 loss 下降但验证 loss 上升 = 过拟合。两者都不下降 = 学习率可能太大或太小。loss 下降突然变慢 = 可能需要调整学习率调度。这些诊断技能在实际训练中非常实用。
- 预训练的成本:即使是 124M 的小模型,在真实规模的数据上预训练也需要大量 GPU 资源。这就是为什么实际工作中很少从头预训练 — 而是加载已有的预训练权重,在此基础上微调。
ℹ️ 总结 — 第5章预训练全流程:文本数据 → 滑动窗口生成 (input, target) 对 → DataLoader 批量加载 → Forward 计算交叉熵损失 → Backward 计算梯度 → AdamW 更新权重(学习率按预热+余弦衰减调度)→ 定期评估 + 生成示例文本。训练完成后,模型学会了语言的统计规律 — 下一步就是通过微调让它学会执行特定任务。
本文是 Build a Large Language Model From Scratch (Sebastian Raschka) 的学习笔记。所有配图版权归原作者所有。代码基于原书示例,有简化和中文注释。