第 9 章 进阶拓展方向

第 9 章 进阶拓展方向

把本项目当起点，往这些方向深入。每节给出最小可行实现思路。

9.1 手写 Transformer

用框架现成的 GPT2 之外，自己实现一遍最能加深理解。分 5 步：

Step 1：单头注意力

import torch
import torch.nn as nn
import torch.nn.functional as F

class SingleHeadAttention(nn.Module):
    def __init__(self, n_embd, head_dim):
        super().__init__()
        self.q = nn.Linear(n_embd, head_dim)
        self.k = nn.Linear(n_embd, head_dim)
        self.v = nn.Linear(n_embd, head_dim)
        self.head_dim = head_dim

    def forward(self, x):
        # x: (B, T, n_embd)
        B, T, _ = x.shape
        q, k, v = self.q(x), self.k(x), self.v(x)
        # 注意力分数 + 缩放
        scores = q @ k.transpose(-2, -1) / (self.head_dim ** 0.5)  # (B, T, T)
        # 因果掩码：只看过去
        mask = torch.tril(torch.ones(T, T, device=x.device))
        scores = scores.masked_fill(mask == 0, float("-inf"))
        attn = F.softmax(scores, dim=-1)
        return attn @ v   # (B, T, head_dim)
​

Step 2：多头注意力

把 N 个单头并行算，结果拼接后过一层投影：

class MultiHeadAttention(nn.Module):
    def __init__(self, n_embd, n_head):
        super().__init__()
        self.heads = nn.ModuleList([
            SingleHeadAttention(n_embd, n_embd // n_head) for _ in range(n_head)
        ])
        self.proj = nn.Linear(n_embd, n_embd)

    def forward(self, x):
        out = torch.cat([h(x) for h in self.heads], dim=-1)
        return self.proj(out)
​

Step 3：加 LayerNorm + 残差，组成 Block

class Block(nn.Module):
    def __init__(self, n_embd, n_head):
        super().__init__()
        self.ln1 = nn.LayerNorm(n_embd)
        self.attn = MultiHeadAttention(n_embd, n_head)
        self.ln2 = nn.LayerNorm(n_embd)
        self.ffn = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd), nn.GELU(),
            nn.Linear(4 * n_embd, n_embd),
        )

    def forward(self, x):
        x = x + self.attn(self.ln1(x))   # 残差 + 注意力
        x = x + self.ffn(self.ln2(x))    # 残差 + FFN
        return x
​

Step 4 + 5：堆叠 N 层 + 嵌入，组成完整 GPT

把 Block 堆 n_layer 层，前面接词嵌入 + 位置嵌入，后面接 LayerNorm + lm_head，就是一个完整的 GPT。替换本项目模型层，用自己写的 GPT 跑训练，对比 loss 下降是否一致。

💡 参考资料：nanoGPT（Karpathy 的 300 行 GPT 实现，本项目精神来源）；Attention Is All You Need（Transformer 原论文）；The Annotated Transformer（逐行注释实现）。

9.2 加验证集与早停

本项目只有训练集，看不到泛化性能。改进思路：

数据切分与验证函数

@torch.no_grad()
def evaluate(model, val_loader, device):
    model.eval()
    total_loss, n = 0, 0
    for x, y in val_loader:
        x, y = x.to(device), y.to(device)
        loss = model(input_ids=x, labels=y).loss
        total_loss += loss.item() * len(x)
        n += len(x)
    model.train()   # 切回训练模式
    return total_loss / n
​

早停

best_val_loss, patience, no_improve = float("inf"), 5, 0

for step in range(max_iters):
    # 训练一步...
    if step % eval_iter == 0:
        val_loss = evaluate(model, val_loader, device)
        if val_loss < best_val_loss:
            best_val_loss, no_improve = val_loss, 0
            save_checkpoint(..., path="checkpoints/best.pt")
        else:
            no_improve += 1
            if no_improve >= patience:
                print("Early stopping!"); break
​

⚠️ 注意陷阱：evaluate 后必须 model.train()；验证要 @torch.no_grad()；不要在验证集上调参，否则信息泄露。

9.3 KV Cache 加速推理

朴素生成的复杂度是 O(seq²)——每生成一个 token，把整段历史重新前向一遍。可以优化到 O(seq)。

朴素：每步重算 T 个位置的 K, V    → 总复杂度 O(T²)
KV Cache：每步只算新位置的 K, V   → 总复杂度 O(T)
​

框架已内置：

out = model.generate(input_ids=..., max_new_tokens=100, use_cache=True)  # 默认就是 True
​

代价：显存占用增加，要存所有层的 K/V。详见《第 6 章》的性能问题讨论。

9.4 更丰富的采样

Top-p（Nucleus）采样

Top-K 固定数量，Top-p 固定概率累加上限：

def sample_top_p(logits, p=0.9):
    probs = F.softmax(logits, dim=-1)
    sorted_probs, sorted_idx = torch.sort(probs, descending=True)
    cumsum = torch.cumsum(sorted_probs, dim=-1)

    # 找累加到 p 的位置
    sorted_indices_to_remove = cumsum > p
    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
    sorted_indices_to_remove[..., 0] = False

    indices_to_remove = sorted_idx[sorted_indices_to_remove]
    probs[indices_to_remove] = 0
    probs = probs / probs.sum()
    return torch.multinomial(probs, num_samples=1)
​

Top-p 的优势：自适应。分布尖锐时只选 3-5 个；分布平坦时可能选几十个。比固定 K=40 更聪明。

其他策略

9.5 混合精度训练（AMP）

scaler = torch.cuda.amp.GradScaler()

for step in range(max_iters):
    x, y = next(...)
    optimizer.zero_grad(set_to_none=True)

    # 前向用 autocast（自动混合精度）
    with torch.cuda.amp.autocast():
        loss = model(input_ids=x, labels=y).loss

    # 反向用 scaler（防 fp16 梯度下溢）
    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)   # 梯度裁剪前先 unscale
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    scaler.step(optimizer)
    scaler.update()
    scheduler.step()
​

收益：显存减半、速度提升 50-100%。梯度缩放防下溢：fp16 能表示的最小正数约 6e-8，小梯度会变 0，scaler 先放大 loss 反向传播时梯度跟着放大，更新前再缩回来。

bf16 vs fp16

新卡（A100/RTX 30+）优先 bf16，溢出风险几乎无，不需要 scaler：

with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
    loss = model(...)
loss.backward()   # 不需要 scaler
​

9.6 更大数据集与分布式训练

数据集升级路径

单机多卡：DDP

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

dist.init_process_group("nccl")
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)

model = model.to(local_rank)
model = DDP(model, device_ids=[local_rank])
# 之后 loss.backward() 自动同步梯度
​

启动：torchrun --nproc_per_node=4 train.py（4 卡并行，每卡处理 1/4 数据，梯度自动平均）。

多机多卡：DeepSpeed ZeRO

9.7 LoRA / 微调预训练模型

不从头训练，加载 gpt2 预训练权重，只训练插入的低秩矩阵。原理：预训练权重 W 不动，旁边加一个低秩更新 BA：

原：y = Wx
LoRA：y = Wx + BAx   （B 是 d×r，A 是 r×d，r << d）
​

只训练 A 和 B，参数量从 d² 降到 2dr。

from peft import LoraConfig, get_peft_model

model = GPT2LMHeadModel.from_pretrained("gpt2")
lora_config = LoraConfig(
    r=8, lora_alpha=16,
    target_modules=["c_attn", "c_proj"],   # 对哪些层加 LoRA
    lora_dropout=0.05,
)
model = get_peft_model(model, lora_config)
# trainable params: 0.24% || all params: 124M
​

只训 0.24% 的参数，效果接近全参数微调，显存骤降。

9.8 Tokenizer 自定义

想用中文或专门领域词表，可用 HuggingFace tokenizers 训练 BPE。关键约束：换了 tokenizer 就必须重新训练 embedding（vocab_size 变了），预训练权重不能直接复用。想微调预训练模型 + 自定义词表需要 model.resize_token_embeddings(new_vocab_size)。

9.9 评估指标

# Loss 越低越好
# Perplexity = exp(loss)，越低越好
import math
ppl = math.exp(loss)
​

perplexity 直觉：「模型平均在每个位置纠结于多少个候选」。PPL=1 完全确定，PPL=50 在 50 个候选间纠结。生成任务还可用 BLEU（n-gram 重叠，翻译/摘要）、ROUGE（召回为主，摘要）；但 LM 自动指标与人工评估相关性不强，最终还是要人评。

9.10 学习路线推荐

必读论文清单

9.11 动手实验

1. 手写 Attention：按 9.1 实现 SingleHeadAttention，喂一个 (1, 4, 8) 的随机张量，看输出 shape 是否正确。
2. 加验证集：按 9.2 改造数据层和训练循环，画 train_loss vs val_loss 曲线，观察是否过拟合。
3. Top-p 实现：按 9.4 在 generate 里加一个 top_p 参数。
4. AMP 改造：按 9.5 给训练循环加混合精度，对比显存和速度。
5. LoRA 微调：按 9.7 用 peft 加载 gpt2，在莎士比亚上微调 100 步，看 loss 是否比从零训降得快。

9.12 结语

恭喜你读完整个教程！🎉

记住：理解一个小项目的每一行，比囫囵跑通一个大项目价值更高。本项目就是为前者设计的。现在去：改一个真实问题（用本项目代码训练你感兴趣的领域文本）、手写一遍（参照 nanoGPT 自己实现注意力）、读一篇论文（Attention Is All You Need 是起点）、上一个台阶（学 HuggingFace Trainer、Accelerate、DeepSpeed）。

祝玩得开心！🚀