第5章：在无标注数据上预训练

学习目标

写出语言模型的交叉熵损失，理解为什么它就是”困惑度”的对数。
实现一个完整的训练循环：前向、损失、反向、优化、评估。
掌握采样阶段的温度、top-k、top-p 三种主流策略。
知道如何加载 OpenAI 公开的 GPT-2 权重。

5.1 损失函数：next-token 交叉熵

预训练的目标极其简洁：让模型对序列的下一个 token 给出尽可能高的概率。形式化为最大似然：

$$\mathcal{L}(\theta) = -\frac{1}{B \cdot T}\sum_{b=1}^{B}\sum_{t=1}^{T} \log P_\theta(x_{t+1}^{(b)} \mid x_{1:t}^{(b)})$$

PyTorch 实现：

import torch
import torch.nn.functional as F

def loss_batch(model, x, y):
    """
    x: (B, T)  输入 ids
    y: (B, T)  目标 ids（输入右移一位）
    """
    logits = model(x)                                  # (B, T, V)
    loss = F.cross_entropy(
        logits.flatten(0, 1),                          # (B*T, V)
        y.flatten(0, 1),                               # (B*T,)
    )
    return loss

困惑度（Perplexity）= $\exp(\mathcal{L})$，可以理解为”模型平均在多少个候选 token 之间犹豫”。完全随机猜的困惑度等于词表大小 $V$，越低越好。

5.2 训练循环骨架

import time

def train(model, train_loader, val_loader, optimizer, device,
          num_epochs=1, eval_freq=200, eval_iter=20):
    train_losses, val_losses = [], []
    seen = 0

    for epoch in range(num_epochs):
        model.train()
        for step, (x, y) in enumerate(train_loader):
            x, y = x.to(device), y.to(device)

            optimizer.zero_grad()
            loss = loss_batch(model, x, y)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)  # 防梯度爆炸
            optimizer.step()

            seen += x.numel()
            if step % eval_freq == 0:
                tr = evaluate(model, train_loader, device, eval_iter)
                va = evaluate(model, val_loader,   device, eval_iter)
                train_losses.append(tr); val_losses.append(va)
                print(f"step {step:5d} | train {tr:.3f} | val {va:.3f}")

    return train_losses, val_losses


@torch.no_grad()
def evaluate(model, loader, device, n_batches):
    model.eval()
    total = 0.0
    for i, (x, y) in enumerate(loader):
        if i >= n_batches: break
        x, y = x.to(device), y.to(device)
        total += loss_batch(model, x, y).item()
    model.train()
    return total / max(n_batches, 1)

几个新手常踩的坑：

忘记 optimizer.zero_grad()——梯度会在 batch 间累加，loss 看起来”很好”但实际是错的；
忘记 .to(device)——把 GPU 模型和 CPU 数据混用会报错；
eval 时忘记 @torch.no_grad()——白白占显存；
eval 时忘记 model.eval()——dropout 仍在生效，损失会偏高。

5.3 优化器与超参数

GPT-2 / GPT-3 系列都使用 AdamW：

optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=3e-4,
    weight_decay=0.1,
    betas=(0.9, 0.95),
)

经验值：

lr 在 $1\text{e-}4$ 到 $6\text{e-}4$ 之间，越大越省训练时间但越容易炸；
weight_decay 0.1，但不要给 LayerNorm 参数和 bias 加 decay（参考附录 D）；
batch tokens 越大越稳，主流做法是用梯度累积凑出 0.5M~4M tokens/step。

5.4 文本生成：从 logits 到字符串

训练好的模型本质上是一个”概率发生器”。生成一段新文本，需要反复执行：

当前序列 ──模型──► logits[-1] ──采样策略──► 下一个 id ──追加到序列 ──循环

最朴素的版本（贪心解码）：

@torch.no_grad()
def generate_greedy(model, idx, max_new_tokens, context_size):
    model.eval()
    for _ in range(max_new_tokens):
        idx_cond = idx[:, -context_size:]          # 截断到上下文窗口
        logits = model(idx_cond)[:, -1, :]         # 只看最后一个时间步: (B, V)
        next_id = logits.argmax(dim=-1, keepdim=True)
        idx = torch.cat([idx, next_id], dim=-1)
    return idx

贪心解码的输出极其单调，所以实践中我们用三种”加噪”手段：

5.4.1 温度采样

把 logits 除以温度 $\tau$ 后再 softmax：

$$P(x) = \mathrm{softmax}(\text{logits} / \tau)$$

$\tau \to 0$：等价于贪心；
$\tau = 1$：原始分布；
$\tau > 1$：分布更平、更”有创意”，但更容易胡说。

5.4.2 Top-k 采样

只在概率最高的 k 个 token 中采样，把其余 logits 设为 $-\infty$。常用 k=40 或 50。

5.4.3 Top-p（Nucleus）采样

按概率从大到小累加，直到累计概率超过阈值 p（如 0.9），只在这个集合内采样。优点是集合大小自适应：分布尖锐时集合小，分布平坦时集合大。

def sample_with_temperature_topk(logits, temperature=1.0, top_k=None):
    if top_k is not None:
        v, _ = torch.topk(logits, top_k)
        logits = torch.where(logits < v[..., -1, None],
                             torch.full_like(logits, float("-inf")),
                             logits)
    probs = torch.softmax(logits / temperature, dim=-1)
    return torch.multinomial(probs, num_samples=1)

5.5 加载官方 GPT-2 权重

OpenAI 在 GPT-2 的官方 TF 实现里发布了 124M / 355M / 774M / 1558M 四个尺寸的权重。原仓库提供了 gpt_download.py，把 TF checkpoint 转换成 numpy 数组后，按字段映射到我们写的 GPTModel。

加载完成后，可以直接：

out = generate(model,
               tokenizer.encode("Once upon a time"),
               max_new_tokens=50,
               temperature=0.8,
               top_k=40)
print(tokenizer.decode(out))

如果你的实现正确，输出会出现”语义连贯”的英文段落——这就是预训练赋予模型的能力。

5.6 关于”完整复现 GPT-2 预训练”的现实

完全从零预训练一个 124M 模型，在主流公开语料（OpenWebText 等）上需要：

~8B tokens；
单张 A100 上约 1 周；
8 卡 A100 约 1 天。

这超出了大多数读者的硬件预算。课程中我们只用一两本古登堡公版书做”形式上的训练”（约 100k tokens），目标是验证训练流程跑得通，而不是得到一个有用的模型。要拿到能生成连贯文本的模型，请直接加载官方权重。

检查清单

我能在不查文档的情况下写出 next-token 交叉熵的实现。
我理解贪心、温度、top-k、top-p 四种解码策略的差异。
我能把 OpenAI 的 TF 权重映射到自己的 GPT 实现。

练习题

在同一段提示上比较温度 $\tau \in {0.2, 0.7, 1.2}$ 的输出，写下你对”创意 vs 一致性”权衡的观察。
自己实现 top-p 采样（不要用现成库），并验证当 $p=1.0$ 时它退化为温度采样。
给训练循环加一个”余弦学习率衰减 + 线性热身”的 scheduler，对比 loss 曲线（参考附录 D）。

📖 第5章补充材料 → — 学习率调度、GPT权重加载、GPT→Llama架构转换

← 上一章 · 返回目录 · 下一章 · 分类微调 →