第 9 章 进阶拓展方向

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

9.1 手写 Transformer

用 HF 现成的 GPT2 之外，自己实现一遍最能加深理解。

从最小可行版本开始

不要一上来就写完整 Transformer。分 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)
        q = self.q(x)   # (B, T, head_dim)
        k = self.k(x)
        v = self.v(x)
        # 注意力分数
        scores = q @ k.transpose(-2, -1) / (self.head_dim ** 0.5)  # (B, T, T)
        # 因果掩码
        mask = torch.tril(torch.ones(T, T))
        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 + 残差

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：堆叠 N 层 + 嵌入

class GPT(nn.Module):
    def __init__(self, vocab_size, n_embd, n_head, n_layer, block_size):
        super().__init__()
        self.wte = nn.Embedding(vocab_size, n_embd)
        self.wpe = nn.Embedding(block_size, n_embd)
        self.blocks = nn.ModuleList([Block(n_embd, n_head) for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size, bias=False)

    def forward(self, idx):
        B, T = idx.shape
        pos = torch.arange(T).unsqueeze(0)
        x = self.wte(idx) + self.wpe(pos)
        for block in self.blocks:
            x = block(x)
        x = self.ln_f(x)
        return self.lm_head(x)   # (B, T, vocab_size)
​

Step 5：测试

替换本项目的 model.py，用自己写的 GPT 跑训练，对比 loss 下降是否一致。

参考资料

• nanoGPT（Karpathy）：GitHub 上 karpathy/nanoGPT 仓库 — 本项目的精神来源，300 行实现 GPT。
• Attention Is All You Need：Transformer 原论文。
• The Annotated Transformer：Harvard NLP 的逐行注释实现。

9.2 加验证集与早停

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

数据切分

# 在 dataset.py 里加
def split_text(text, val_ratio=0.1):
    split = int(len(text) * (1 - val_ratio))
    return text[:split], text[split:]
​

或更合理的：按 token 切分。

验证函数

@torch.no_grad()
def evaluate(model, val_loader, device):
    model.eval()
    total_loss = 0
    n = 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 = float("inf")
patience = 5  # 验证 loss 连续 5 次不降就停
no_improve = 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 = val_loss
            no_improve = 0
            save_checkpoint(..., path="checkpoints/best.pt")
        else:
            no_improve += 1
            if no_improve >= patience:
                print("Early stopping!")
                break
​

注意陷阱

• 切回 train 模式：evaluate 后必须 model.train()，否则后续训练 dropout 关闭。
• @torch.no_grad()：验证不建图，省显存。
• 不要在验证集上调参：否则信息泄露，val_loss 不可信。需要的话再切一个 test set。

9.3 KV Cache 加速推理

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

原理

注意力计算里，每个历史位置的 K 和 V 是不变的（只是越来越多）。把它们缓存起来，下一步只算新 token 的 Q 与所有历史 K/V 做注意力。

朴素：每步重算 T 个位置的 K, V
       step 1: 算 K, V for [t0]           → 1 次前向
       step 2: 算 K, V for [t0, t1]       → 2 次前向
       step 3: 算 K, V for [t0, t1, t2]   → 3 次前向
       ...
       总复杂度 O(T²)

KV Cache：每步只算新位置的 K, V
       step 1: 算 K, V for t0, 缓存       → 1 次前向
       step 2: 算 K, V for t1, 追加缓存   → 1 次前向
       step 3: 算 K, V for t2, 追加缓存   → 1 次前向
       ...
       总复杂度 O(T)
​

HF 已内置

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

或在手写循环里：

past_key_values = None
for _ in range(max_new_tokens):
    out = model(input_ids=next_token, past_key_values=past_key_values)
    past_key_values = out.past_key_values   # 缓存
    next_token = sample(out.logits[:, -1])
​

代价

显存占用增加：要存所有层的 K/V（shape (B, n_head, T, head_dim) × 2 × n_layer）。长序列时这是大头。

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 的优势：自适应。

• 当分布尖锐（模型很自信），Top-p=0.9 可能只选 3-5 个 token。
• 当分布平坦（模型迷茫），Top-p=0.9 可能选几十个。

比固定 K=40 更聪明。

Typical Sampling

选「信息熵接近平均」的 token，避免极端高概率（重复）和极端低概率（乱来）。论文 Hierarchical Neural Story Generation 提出。

Contrastive Search

新方法（2022），用一个罚项平衡连贯性和多样性。HF 已支持：

out = model.generate(
    input_ids=...,
    do_sample=False,
    penalty_alpha=0.6,
    top_k=4,
)
​

Beam Search

维护 K 条候选路径，每步扩展，选整体概率最大的。适合翻译/摘要（有标准答案），不适合开放生成（容易平庸重复）。

out = model.generate(
    input_ids=...,
    num_beams=4,
    max_new_tokens=100,
)
​

9.5 混合精度训练（AMP）

fp32 训练显存占用大。AMP 自动在 fp16/bf16 和 fp32 间切换，省显存、提速度。

完整改造

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()
    
    # 梯度裁剪前先 unscale（否则裁剪的是放大后的梯度）
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    
    # 步进
    scaler.step(optimizer)
    scaler.update()
    scheduler.step()
​

收益

• 显存减半（fp16 占 2 字节 vs fp32 4 字节）。
• 速度提升 50-100%（Tensor Core 加速 fp16 矩阵乘）。
• 梯度缩放防下溢：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 更大数据集与分布式训练

数据集升级路径

流式加载

大数据集一次性载入内存不现实，用流式：

from datasets import load_dataset
ds = load_dataset("openwebtext", split="train", streaming=True)
for example in ds:
    text = example["text"]
    # 流式处理，不全载入
​

单机多卡：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])

sampler = DistributedSampler(dataset)
dataloader = DataLoader(dataset, sampler=sampler, ...)

for x, y in dataloader:
    loss = model(x, labels=y).loss
    loss.backward()   # DDP 自动同步梯度
​

启动：

torchrun --nproc_per_node=4 train.py
# 4 卡并行，每卡处理 1/4 数据，梯度自动平均
​

多机多卡：DeepSpeed ZeRO

DDP 每卡存完整模型副本，显存吃不消。ZeRO 把权重/优化器/梯度切到多卡：

import deepspeed
model, optimizer, _, _ = deepspeed.initialize(
    model=model, optimizer=optimizer, config=ds_config
)
​

9.7 LoRA / 微调预训练模型

不从头训练，加载 gpt2 预训练权重，只训练插入的低秩矩阵。

原理

预训练权重 W 不动，旁边加一个低秩更新 BA：

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

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

用 PEFT 库

from peft import LoraConfig, get_peft_model
from transformers import GPT2LMHeadModel

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)
model.print_trainable_parameters()
# trainable params: 294,912 || all params: 124,734,720 || trainable%: 0.24%
​

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

9.8 Tokenizer 自定义

想用中文或专门领域词表：

用 HuggingFace tokenizers 训练 BPE

from tokenizers import Tokenizer
from tokenizers.models import BPE
from tokenizers.trainers import BpeTrainer
from tokenizers.pre_tokenizers import Whitespace

tokenizer = Tokenizer(BPE(unk_token="<unk>"))
tokenizer.pre_tokenizer = Whitespace()
trainer = BpeTrainer(
    vocab_size=30000,
    special_tokens=["<unk>", "<pad>", "<bos>", "<eos>"],
)
tokenizer.train(files=["corpus.txt"], trainer=trainer)
​

关键约束

换了 tokenizer 就必须重新训练 embedding（vocab_size 变了）：

# 旧：vocab_size=50257 (p50k_base)
# 新：vocab_size=30000 (自定义)
# 模型的 wte 和 lm_head 都要重建
​

预训练权重不能直接复用（词表对不上）。如果想微调预训练模型 + 自定义词表，需要 resize：

model = GPT2LMHeadModel.from_pretrained("gpt2")
model.resize_token_embeddings(new_vocab_size)
​

9.9 评估指标

Loss / Perplexity

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

perplexity 直觉：「模型平均在每个位置纠结于多少个候选」。PPL=1 表示完全确定，PPL=50 表示在 50 个候选间纠结。

BLEU / ROUGE（生成任务）

• BLEU：n-gram 重叠，翻译/摘要常用。
• ROUGE：召回为主，摘要常用。

人工评估

• 流畅度、连贯性、相关性。
• LM 自动指标与人工评估相关性不强，最终还是要人评。

9.10 学习路线推荐

阶段 1：跑通本项目（1-2 周）

• 装环境，跑通 train + inference + app。
• 改超参，观察 loss / 生成质量变化。

阶段 2：手写 Transformer（2-4 周）

• 跟着 nanoGPT 实现一遍。
• 对比与 HF GPT2 的差异。

阶段 3：微调预训练模型（1-2 周）

• 加载 gpt2，在自己的数据上微调。
• 学 LoRA、Prompt Tuning。

阶段 4：训练更大模型（按需）

• 用 DDP / DeepSpeed 训 100M+ 模型。
• 学习分布式调试技巧。

阶段 5：研究前沿

• 读论文：LLaMA、GPT-4 技术报告、InstructGPT（RLHF）、DPO 等。
• 跟进 HuggingFace accelerate、trl、transformers 新功能。

必读论文清单

9.11 动手实验

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

9.12 小结

进阶方向 10 个：

1. 手写 Transformer（从单头注意力开始）
2. 验证集 + 早停
3. KV Cache 加速推理
4. Top-p / Typical / Contrastive 采样
5. AMP 混合精度训练
6. DDP / DeepSpeed 分布式
7. LoRA 参数高效微调
8. 自定义 Tokenizer
9. 评估指标（PPL / BLEU / 人工）
10. 论文阅读路线

每条都是独立可深入的领域，选感兴趣的钻。

9.13 结语

恭喜你读完整个教程！🎉

记住：理解一个小项目的每一行，比囫囵跑通一个大项目价值更高。本项目就是为前者设计的。

现在去：

1. 改一个真实问题：用本项目代码做底座，训练你感兴趣的领域文本（古诗、代码、对话）。
2. 手写一遍：参照 nanoGPT，自己实现 MultiHeadAttention。
3. 读一篇论文：Attention Is All You Need 是起点。
4. 上一个台阶：学 HuggingFace Trainer、Accelerate、DeepSpeed。

祝玩得开心！🚀