从零构建LLM 第5章:预训练

📅 2026-03-02📖 ~9 min readLLMPretrainingTraining
本文是 Build a Large Language Model From Scratch(Sebastian Raschka 著)第5章的学习笔记。上一章搭建了 GPT 的完整架构,但权重还是随机的 — 模型不会说话。本章的核心任务是通过预训练 (Pretraining) 让模型学会语言:用文本数据,通过下一个 token 预测这个看似简单的任务,训练模型内化语言的统计规律。同时介绍了文本生成策略(temperature / top-k)和学习率调度等关键训练技巧。
← Ch4: GPT架构Ch6: 文本分类微调 →

1. 预训练的目标:下一个 Token 预测

LLM 的预训练目标极其简洁:给定前 t 个 token,预测第 t+1 个 token。这就是第2章中滑动窗口生成的 (input, target) 对的用武之地。

Chapter 5 overview

第5章在整体 LLM 构建流程中的位置 — 预训练是让模型「学会语言」的关键步骤 (图源: LLMs from Scratch)

衡量预测好坏的指标是交叉熵损失 (Cross-Entropy Loss):模型输出一个概率分布(50257 维向量),真实答案是一个 token ID。交叉熵衡量的是模型的概率分布与真实分布之间的「距离」— 损失越小,说明模型给正确 token 分配的概率越高。

Python — 损失计算
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 的标配。

Training and validation loss

训练过程中的 loss 曲线 — 随着训练进行,模型逐渐学会更好地预测下一个 token (图源: LLMs from Scratch)
Python — 训练循环 (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 = [], [], []
    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 中采样
Temperature effect on generation

温度对生成结果的影响:低温 → 保守确定,高温 → 多样随机 (图源: LLMs from Scratch)
Python — 文本生成 (temperature + top-k)
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 普遍采用的线性预热 + 余弦衰减策略:

  1. 预热阶段 (Warmup):学习率从 0 线性增长到峰值。避免训练初期步子太大破坏随机初始化的权重分布
  2. 余弦衰减 (Cosine Decay):学习率按余弦曲线从峰值平滑下降到最小值。后期用小学习率精细调优
Learning rate schedule

线性预热 + 余弦衰减的学习率曲线 (图源: LLMs from Scratch)
Python — 学习率调度
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) 的学习笔记。所有配图版权归原作者所有。代码基于原书示例,有简化和中文注释。

← Ch4: GPT架构Ch6: 文本分类微调 →

Related Posts