第21章训练一个简单的 GPT

本章目标：

去掉 Encoder 和 Cross-Attention，只用 Decoder，在莎士比亚全集上从零训练一个 1000 万参数的字符级 GPT；用 SDPA + bf16 + torch.compile 在 RTX 4090 上 90 秒训完；观察过拟合并用最佳 checkpoint 早停；实现 Temperature / Top-k 采样；亲手实现 KV Cache，验证它与朴素生成逐字一致，并找出它真正发挥作用的条件。

本章目录

21.1 本章做什么 ------ 字符级语言模型、与第14、20章代码的区别
21.2 数据准备 ------ 字符级 Tokenizer、数据整体放进显存
21.3 模型 ------ 基于 SDPA 的 Causal Self-Attention、GPT Block、GPT
21.4 采样与生成 ------ Greedy / Temperature / Top-k
21.5 训练 ------ Warmup + Cosine、AdamW、梯度裁剪、GPU 加速四件套
21.6 运行结果 ------ 过拟合曲线与早停
21.7 实验：Temperature 对生成的影响
21.8 实现 KV Cache ------ Prefill / Decode、一致性验证
21.9 实验：KV Cache 什么时候才有用？
21.10 Hugging Face 源码定位
21.11 常见坑

21.1 本章做什么

21.1.1 为什么从字符级开始？

方案	词表大小	优点	缺点
字符级	65	无需训练 Tokenizer，代码最短；loss 可以和 ln⁡65\ln 65ln65 直接对比	序列长，同样的上下文窗口能看到的内容少
BPE（GPT-2）	50257	序列短、效率高	`d_model=384` 时，光 Embedding（与输出层共享）就要 1930 万参数，比本章整个模型还大

第14章 14.10 节讲的 BPE 可以作为练习替换进来（见章末）。

21.1.2 与前面代码的区别

	第14章 14.11	第20章	本章
结构	Decoder-Only	Encoder-Decoder	Decoder-Only
Attention	Keras 内置	手写（要画注意力图）	SDPA（FlashAttention 内核）
Q/K/V 投影	三个矩阵	三个矩阵	一个 `Linear(d, 3d)`（GPT-2 的 `c_attn`）
推理	朴素生成	朴素生成	朴素生成 + KV Cache
规模	示意	93 万参数	1077 万参数

21.1.3 整体结构

#mermaid-svg-YuHxW0NzGqRJg1Uf{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-YuHxW0NzGqRJg1Uf .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-YuHxW0NzGqRJg1Uf .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-YuHxW0NzGqRJg1Uf .error-icon{fill:#552222;}#mermaid-svg-YuHxW0NzGqRJg1Uf .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-YuHxW0NzGqRJg1Uf .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-YuHxW0NzGqRJg1Uf .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-YuHxW0NzGqRJg1Uf .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-YuHxW0NzGqRJg1Uf .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-YuHxW0NzGqRJg1Uf .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-YuHxW0NzGqRJg1Uf .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-YuHxW0NzGqRJg1Uf .marker{fill:#333333;stroke:#333333;}#mermaid-svg-YuHxW0NzGqRJg1Uf .marker.cross{stroke:#333333;}#mermaid-svg-YuHxW0NzGqRJg1Uf svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-YuHxW0NzGqRJg1Uf p{margin:0;}#mermaid-svg-YuHxW0NzGqRJg1Uf .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-YuHxW0NzGqRJg1Uf .cluster-label text{fill:#333;}#mermaid-svg-YuHxW0NzGqRJg1Uf .cluster-label span{color:#333;}#mermaid-svg-YuHxW0NzGqRJg1Uf .cluster-label span p{background-color:transparent;}#mermaid-svg-YuHxW0NzGqRJg1Uf .label text,#mermaid-svg-YuHxW0NzGqRJg1Uf span{fill:#333;color:#333;}#mermaid-svg-YuHxW0NzGqRJg1Uf .node rect,#mermaid-svg-YuHxW0NzGqRJg1Uf .node circle,#mermaid-svg-YuHxW0NzGqRJg1Uf .node ellipse,#mermaid-svg-YuHxW0NzGqRJg1Uf .node polygon,#mermaid-svg-YuHxW0NzGqRJg1Uf .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-YuHxW0NzGqRJg1Uf .rough-node .label text,#mermaid-svg-YuHxW0NzGqRJg1Uf .node .label text,#mermaid-svg-YuHxW0NzGqRJg1Uf .image-shape .label,#mermaid-svg-YuHxW0NzGqRJg1Uf .icon-shape .label{text-anchor:middle;}#mermaid-svg-YuHxW0NzGqRJg1Uf .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-YuHxW0NzGqRJg1Uf .rough-node .label,#mermaid-svg-YuHxW0NzGqRJg1Uf .node .label,#mermaid-svg-YuHxW0NzGqRJg1Uf .image-shape .label,#mermaid-svg-YuHxW0NzGqRJg1Uf .icon-shape .label{text-align:center;}#mermaid-svg-YuHxW0NzGqRJg1Uf .node.clickable{cursor:pointer;}#mermaid-svg-YuHxW0NzGqRJg1Uf .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-YuHxW0NzGqRJg1Uf .arrowheadPath{fill:#333333;}#mermaid-svg-YuHxW0NzGqRJg1Uf .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-YuHxW0NzGqRJg1Uf .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-YuHxW0NzGqRJg1Uf .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-YuHxW0NzGqRJg1Uf .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-YuHxW0NzGqRJg1Uf .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-YuHxW0NzGqRJg1Uf .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-YuHxW0NzGqRJg1Uf .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-YuHxW0NzGqRJg1Uf .cluster text{fill:#333;}#mermaid-svg-YuHxW0NzGqRJg1Uf .cluster span{color:#333;}#mermaid-svg-YuHxW0NzGqRJg1Uf div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-YuHxW0NzGqRJg1Uf .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-YuHxW0NzGqRJg1Uf rect.text{fill:none;stroke-width:0;}#mermaid-svg-YuHxW0NzGqRJg1Uf .icon-shape,#mermaid-svg-YuHxW0NzGqRJg1Uf .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-YuHxW0NzGqRJg1Uf .icon-shape p,#mermaid-svg-YuHxW0NzGqRJg1Uf .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-YuHxW0NzGqRJg1Uf .icon-shape .label rect,#mermaid-svg-YuHxW0NzGqRJg1Uf .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-YuHxW0NzGqRJg1Uf .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-YuHxW0NzGqRJg1Uf .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-YuHxW0NzGqRJg1Uf :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} 字符序列 idx

B, T

Token Embedding

可学习位置 Embedding
GPT Block × 6

LN → Causal Self-Attn → +

LN → MLP → +
最终 LayerNorm
lm_head（与 Token Embedding 共享权重）

B, T, 65

Cross-Entropy

与右移一位的标签比较

text 复制代码

ch21_gpt/
├── gpt.py           # 模型 + 采样 + 两种生成函数（21.3、21.4、21.8）
├── train_gpt.py     # 数据、训练、早停、生成（21.2、21.5）
├── common.py        # 加载 checkpoint
├── temp_exp.py      # Temperature 实验（21.7）
├── test_cache.py    # KV Cache 一致性与速度（21.8）
└── bench_cache.py   # KV Cache 在不同 batch / 长度下的加速比（21.9）

环境与第20章相同：Python 3.12、PyTorch 2.7.0+cu126、RTX 4090 24GB。下文所有输出都来自真实运行。

21.2 数据准备

python 复制代码

# ch21_gpt/train_gpt.py
"""第21章：在莎士比亚文本上训练字符级 GPT
 
运行：python train_gpt.py
配置参考 nanoGPT 的 shakespeare_char（6 层、384 维、上下文 256），RTX 4090 上几分钟可跑完。
"""
import math
import os
import time
import urllib.request
 
import torch
 
from gpt import GPT, generate
 
# ---------------- 0. 设备与精度 ----------------
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
USE_BF16 = DEVICE == "cuda" and torch.cuda.is_bf16_supported()
USE_COMPILE = DEVICE == "cuda" and os.name != "nt"      # torch.compile 依赖 Triton，Windows 上默认关闭
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.manual_seed(1337)
 
# ---------------- 1. 数据 ----------------
URL = "https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
if not os.path.exists("input.txt"):
    urllib.request.urlretrieve(URL, "input.txt")
text = open("input.txt", encoding="utf-8").read()
 
chars = sorted(set(text))
vocab_size = len(chars)
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for c, i in stoi.items()}
encode = lambda s: [stoi[c] for c in s]
decode = lambda ids: "".join(itos[int(i)] for i in ids)
 
data = torch.tensor(encode(text), dtype=torch.long)
n = int(0.9 * len(data))
train_data, val_data = data[:n].to(DEVICE), data[n:].to(DEVICE)   # 只有 1MB，整体放进显存
print(f"字符数 {len(data):,}，词表大小 {vocab_size}")

text 复制代码

字符数 1,115,394，词表大小 65

整份数据只有 1MB（110 万个 int64 约 9MB），直接整体放进显存。这样取 batch 时不需要 CPU→GPU 的拷贝，也不需要 DataLoader。

文件开头的"设备与精度"几行在 21.5.3 节统一解释。

训练样本不需要预先切好，每一步随机抽 BATCH_SIZE 个起点，各取一段长 BLOCK_SIZE 的窗口：

python 复制代码

# ch21_gpt/train_gpt.py（续）
# ---------------- 2. 超参数 ----------------
BLOCK_SIZE = 256
BATCH_SIZE = 64
D_MODEL, N_HEAD, N_LAYER, DROPOUT = 384, 6, 6, 0.2
MAX_STEPS, WARMUP = 5000, 100
LR_MAX, LR_MIN = 1e-3, 1e-4
WEIGHT_DECAY, BETAS, GRAD_CLIP = 0.1, (0.9, 0.99), 1.0
EVAL_EVERY, EVAL_ITERS = 250, 100
# CPU 调试时可改小：BLOCK_SIZE=128, BATCH_SIZE=32, D_MODEL=128, N_HEAD=4, N_LAYER=4, MAX_STEPS=3000
 
 
def get_batch(split):
    d = train_data if split == "train" else val_data
    ix = torch.randint(len(d) - BLOCK_SIZE - 1, (BATCH_SIZE,), device=d.device)
    offsets = torch.arange(BLOCK_SIZE, device=d.device)
    x = d[ix[:, None] + offsets]                         # [B, T]，用索引一次取出，不走 Python 循环
    y = d[ix[:, None] + offsets + 1]                     # 右移一位就是标签
    return x, y

d[ix[:, None] + offsets] 是一个索引技巧：ix[:, None] 的 Shape 是 [B, 1]，offsets 是 [T]，相加后广播成 [B, T] 的下标矩阵，一次取出整个 batch，不用 Python 循环。

一个窗口其实包含 256 个训练样本（第14章 14.3.1）：

text 复制代码

x: F  i  r  s  t     C  i  t  ...
y: i  r  s  t     C  i  t  i  ...
   ↑ 位置 0 看到 "F"，预测 "i"
      ↑ 位置 1 看到 "Fi"，预测 "r"

21.3 模型

21.3.1 Causal Self-Attention：用 SDPA，并预留 KV Cache 接口

python 复制代码

# ch21_gpt/gpt.py
"""第21章 字符级 GPT（PyTorch 版）------ 模型、采样、KV Cache 生成
 
与第20章不同，这里的 Attention 用 F.scaled_dot_product_attention（SDPA）：
在 RTX 4090 上会自动选用 FlashAttention / Memory-Efficient 内核。
"""
import math
 
import torch
import torch.nn as nn
import torch.nn.functional as F
 
 
class CausalSelfAttention(nn.Module):
    def __init__(self, d_model, num_heads, dropout=0.1):
        super().__init__()
        assert d_model % num_heads == 0
        self.h, self.d_k = num_heads, d_model // num_heads
        self.qkv = nn.Linear(d_model, 3 * d_model)          # 一次投影出 Q/K/V（GPT-2 的 c_attn）
        self.proj = nn.Linear(d_model, d_model)
        self.dropout = dropout
        self.resid_drop = nn.Dropout(dropout)
 
    def forward(self, x, cache=None):
        B, T, C = x.shape
        q, k, v = self.qkv(x).split(C, dim=-1)
        q, k, v = (t.view(B, T, self.h, self.d_k).transpose(1, 2) for t in (q, k, v))  # [B,h,T,d_k]
 
        past = 0
        if cache is not None and cache["k"] is not None:     # 推理：拼上历史 K/V
            past = cache["k"].size(2)
            k = torch.cat([cache["k"], k], dim=2)
            v = torch.cat([cache["v"], v], dim=2)
        new_cache = {"k": k, "v": v} if cache is not None else None
 
        dropout_p = self.dropout if self.training else 0.0
        if past == 0:
            # 训练 / Prefill：Tq == Tk，标准下三角
            y = F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=dropout_p)
        elif T == 1:
            # Decode 单步：新 Token 可以看到全部历史，不需要任何 Mask
            y = F.scaled_dot_product_attention(q, k, v, dropout_p=dropout_p)
        else:
            # 有历史又一次喂多个 Token（如分块 Prefill）：不能用 is_causal=True！
            # SDPA 的 is_causal 在 Tq != Tk 时按"左上角"对齐，会让新 Token 看不到大部分历史。
            # 这里显式构造"右下角"对齐的 Mask。注意 SDPA 的 bool mask：True = 可见。
            Tk = k.size(2)
            mask = torch.ones(T, Tk, dtype=torch.bool, device=x.device).tril(diagonal=Tk - T)
            y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=dropout_p)
 
        y = y.transpose(1, 2).reshape(B, T, C)
        return self.resid_drop(self.proj(y)), new_cache

第20章 20.14 节已经测过：SDPA 在 4090 上比手写 Attention 快约 12 倍、显存约为 1/9。代价是拿不到注意力权重（FlashAttention 从不把完整的权重矩阵写出来），所以第20章要画图时只能手写。

这段代码最关键的是三个分支。SDPA 的 is_causal=True 有一个容易忽略的细节：当 Query 长度 TqT_qTq 和 Key 长度 TkT_kTk 不相等时，它生成的下三角是按"左上角"对齐的：

text 复制代码

Tq = 2, Tk = 5，is_causal=True 生成的 Mask（1 = 可见）：
            k0 k1 k2 k3 k4
  q0 (新)   1  0  0  0  0      ← 新 Token 只能看到 k0？错！
  q1 (新)   1  1  0  0  0
 
正确的 Mask（按"右下角"对齐）：
            k0 k1 k2 k3 k4
  q0 (新)   1  1  1  1  0      ← 能看到全部 3 个历史 + 自己
  q1 (新)   1  1  1  1  1

带 Cache 时，新 Token 排在 Key 序列的最后面，正确的做法是按右下角对齐。所以分三种情况处理：

场景	TqT_qTq vs TkT_kTk	做法
训练 / Prefill（没有历史）	相等	`is_causal=True`，左上角和右下角对齐是一回事
Decode 单步	Tq=1T_q = 1Tq=1	不需要 Mask：新 Token 本来就能看到全部历史和自己
有历史又一次喂多个 Token（分块 Prefill、投机解码）	Tq<TkT_q < T_kTq<Tk	显式构造右下角对齐的 Mask：`tril(diagonal=Tk-Tq)`

再强调一次 Mask 语义：SDPA 的 bool attn_mask 是 True = 可见 ，与第20章 masked_fill 的 True = 屏蔽正好相反。

21.3.2 GPT Block 与 GPT

python 复制代码

# ch21_gpt/gpt.py（续）
class GPTBlock(nn.Module):
    """Pre-Norm Decoder Block：没有 Cross-Attention"""
    def __init__(self, d_model, num_heads, dropout=0.1):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = CausalSelfAttention(d_model, num_heads, dropout)
        self.ln2 = nn.LayerNorm(d_model)
        self.mlp = nn.Sequential(
            nn.Linear(d_model, 4 * d_model), nn.GELU(),
            nn.Linear(4 * d_model, d_model), nn.Dropout(dropout),
        )
 
    def forward(self, x, cache=None):
        a, new_cache = self.attn(self.ln1(x), cache)
        x = x + a
        x = x + self.mlp(self.ln2(x))
        return x, new_cache
 
 
class GPT(nn.Module):
    def __init__(self, vocab_size, block_size=256, d_model=384,
                 num_heads=6, num_layers=6, dropout=0.2, tie_weights=True):
        super().__init__()
        self.block_size = block_size
        self.tok_emb = nn.Embedding(vocab_size, d_model)
        self.pos_emb = nn.Embedding(block_size, d_model)     # 可学习位置编码
        self.drop = nn.Dropout(dropout)
        self.blocks = nn.ModuleList(GPTBlock(d_model, num_heads, dropout) for _ in range(num_layers))
        self.ln_f = nn.LayerNorm(d_model)
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
        if tie_weights:                                      # Weight Tying（GPT-2 做法）
            self.lm_head.weight = self.tok_emb.weight
        self.apply(self._init_weights)
        # GPT-2：残差分支的输出投影按 1/sqrt(2·层数) 缩放初始化，深层更稳
        for name, p in self.named_parameters():
            if name.endswith("attn.proj.weight") or name.endswith("mlp.2.weight"):
                nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * num_layers))
 
    @staticmethod
    def _init_weights(m):
        if isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, mean=0.0, std=0.02)
            if m.bias is not None:
                nn.init.zeros_(m.bias)
        elif isinstance(m, nn.Embedding):
            nn.init.normal_(m.weight, mean=0.0, std=0.02)
 
    def forward(self, idx, targets=None, caches=None, start_pos=0):
        B, T = idx.shape
        pos = torch.arange(start_pos, start_pos + T, device=idx.device)
        x = self.drop(self.tok_emb(idx) + self.pos_emb(pos))
        new_caches = []
        for i, blk in enumerate(self.blocks):
            x, c = blk(x, None if caches is None else caches[i])
            new_caches.append(c)
        logits = self.lm_head(self.ln_f(x))                  # [B, T, vocab]
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)).float(), targets.view(-1))
        if caches is not None:
            return logits, new_caches
        return logits, loss
 
    def configure_optimizer(self, lr, weight_decay, betas, device_type):
        """二维参数（矩阵、Embedding）做 weight decay；bias 和 LayerNorm 不做"""
        params = [p for p in self.parameters() if p.requires_grad]
        decay = [p for p in params if p.dim() >= 2]
        no_decay = [p for p in params if p.dim() < 2]
        groups = [{"params": decay, "weight_decay": weight_decay},
                  {"params": no_decay, "weight_decay": 0.0}]
        # fused=True：CUDA 上把 AdamW 的逐参数更新融合成少数几个 kernel
        return torch.optim.AdamW(groups, lr=lr, betas=betas, fused=(device_type == "cuda"))

三处 GPT-2 的做法：

Weight Tying ：lm_head.weight = tok_emb.weight，输出层和输入 Embedding 共享同一个矩阵（第14章 14.11.1）。词表只有 65 时省下的参数不多，但这是 GPT-2 / Llama 系列的标准做法。
初始化 ：所有权重 N(0,0.02)\mathcal{N}(0, 0.02)N(0,0.02)；每个残差分支的最后一层（attn.proj、mlp.2）再缩小 1/2N1/\sqrt{2N}1/2N 。每个 Block 往主干上加两次残差，NNN 层共加 2N2N2N 次，缩小初始值可以防止主干的方差随层数线性增长。
Weight Decay 分组：只对二维参数（矩阵、Embedding）做权重衰减，bias 和 LayerNorm 的参数不衰减------对它们做衰减只会把它们往 0 拉，没有正则化意义。

forward 里两个推理专用参数：

caches：每层一个 {"k", "v"} 字典。传入时返回 (logits, new_caches)，不传时返回 (logits, loss)，训练代码不受影响。
start_pos：带 Cache 时每次只喂 1 个新 Token，它的位置编号不是 0，而是它在整段序列中的真实位置。

本章配置（6 层、d_model=384）参数量 10.77M 。每个 Block 约 12d2≈17712 d^2 ≈ 17712d2≈177 万参数（Attention 4d24d^24d2 + MLP 8d28d^28d2，第7、8章），6 层约 1062 万，其余是位置 Embedding（256×384≈9.8256 \times 384 ≈ 9.8256×384≈9.8 万）、Token Embedding（与输出层共享，65×384≈2.565 \times 384 ≈ 2.565×384≈2.5 万）和 bias / LayerNorm。

21.4 采样与生成

python 复制代码

# ch21_gpt/gpt.py（续）
def sample_next(logits, temperature=1.0, top_k=None):
    """logits: [B, vocab] -> [B, 1]"""
    if temperature == 0:                                     # 贪心
        return logits.argmax(-1, keepdim=True)
    logits = logits / temperature
    if top_k is not None:
        kth = torch.topk(logits, min(top_k, logits.size(-1))).values[:, -1:]
        logits = logits.masked_fill(logits < kth, float("-inf"))   # 只保留前 k 个
    return torch.multinomial(F.softmax(logits.float(), dim=-1), num_samples=1)
 
 
@torch.no_grad()
def generate(model, idx, max_new_tokens, temperature=1.0, top_k=None):
    """朴素生成：每一步都把整段上下文重新跑一遍"""
    model.eval()
    for _ in range(max_new_tokens):
        idx_cond = idx[:, -model.block_size:]                # 超过上下文窗口就截断
        logits, _ = model(idx_cond)
        idx = torch.cat([idx, sample_next(logits[:, -1], temperature, top_k)], dim=1)
    return idx

参数	作用（第14章 14.12）
`temperature=0`	贪心，永远选概率最大的字符，结果确定
`temperature<1`	分布更尖锐，更保守
`temperature>1`	分布更平坦，更有"创意"，也更容易胡言乱语
`top_k=20`	只在概率最高的 20 个字符里采样，砍掉长尾噪声

两个 PyTorch 细节：Top-k 用 masked_fill(..., -inf) 把其余位置的概率置为 0；torch.multinomial 前把 logits 转回 FP32，避免 bf16 下 Softmax 精度不足导致概率之和偏离 1。

21.5 训练

21.5.1 学习率：Warmup + Cosine

python 复制代码

# ch21_gpt/train_gpt.py（续）
def lr_at(step):
    """Warmup + Cosine"""
    if step < WARMUP:
        return LR_MAX * (step + 1) / WARMUP
    progress = min(1.0, (step - WARMUP) / (MAX_STEPS - WARMUP))
    return LR_MIN + 0.5 * (LR_MAX - LR_MIN) * (1 + math.cos(math.pi * progress))

前 100 步线性升到 1e-3，之后按余弦曲线平滑退火到 1e-4。训练步数固定时，Cosine 比第20章的反平方根衰减更常用（GPT-3、Llama 都是 Warmup + Cosine）。

21.5.2 训练循环与早停

python 复制代码

# ch21_gpt/train_gpt.py（续）
# ---------------- 3. 模型与优化器 ----------------
raw_model = GPT(vocab_size, BLOCK_SIZE, D_MODEL, N_HEAD, N_LAYER, DROPOUT).to(DEVICE)
print(f"device={DEVICE} bf16={USE_BF16} compile={USE_COMPILE} "
      f"参数量 {sum(p.numel() for p in raw_model.parameters()) / 1e6:.2f}M")
optimizer = raw_model.configure_optimizer(LR_MAX, WEIGHT_DECAY, BETAS, DEVICE)
# 训练用编译后的模型；生成时用 raw_model（序列长度每步都变，编译反而会反复重编译）
model = torch.compile(raw_model) if USE_COMPILE else raw_model
autocast = lambda: torch.autocast(device_type=DEVICE, dtype=torch.bfloat16, enabled=USE_BF16)
 
 
@torch.no_grad()
def estimate_loss():
    model.eval()
    out = {}
    for split in ("train", "val"):
        losses = torch.zeros(EVAL_ITERS)
        for i in range(EVAL_ITERS):
            with autocast():
                _, loss = model(*get_batch(split))
            losses[i] = loss.item()
        out[split] = losses.mean().item()
    model.train()
    return out
 
 
# ---------------- 4. 训练循环 ----------------
print(f"初始 loss 理论值 ln({vocab_size}) = {math.log(vocab_size):.3f}")
CKPT = "gpt_shakespeare.pt"
config = dict(vocab_size=vocab_size, block_size=BLOCK_SIZE, d_model=D_MODEL,
              num_heads=N_HEAD, num_layers=N_LAYER, dropout=DROPOUT)
best_val, best_step = float("inf"), -1
 
model.train()
t0 = time.time()
for step in range(MAX_STEPS + 1):
    if step % EVAL_EVERY == 0:
        l = estimate_loss()
        mark = ""
        # 1MB 的数据对 10M 参数的模型太少，训练后期一定会过拟合：
        # 只在 val loss 创新低时保存（nanoGPT 的 shakespeare_char 配置也是这么做的）
        if l["val"] < best_val:
            best_val, best_step, mark = l["val"], step, "  *saved"
            torch.save({"model": raw_model.state_dict(), "chars": chars,
                        "config": config, "step": step, "val_loss": best_val}, CKPT)
        print(f"step {step:5d} | train {l['train']:.4f} | val {l['val']:.4f} | "
              f"gap {l['val'] - l['train']:.3f} | lr {lr_at(step):.2e} | {time.time() - t0:.0f}s{mark}")
    if step == MAX_STEPS:
        break
 
    for g in optimizer.param_groups:
        g["lr"] = lr_at(step)
    x, y = get_batch("train")
    with autocast():
        _, loss = model(x, y)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(raw_model.parameters(), GRAD_CLIP)
    optimizer.step()
 
if DEVICE == "cuda":
    print(f"峰值显存 {torch.cuda.max_memory_allocated() / 2**30:.2f} GB")
 
print(f"最佳 val loss {best_val:.4f}（step {best_step}），已保存到 {CKPT}")
 
# ---------------- 5. 生成：用最佳 checkpoint，而不是过拟合后的最终权重 ----------------
raw_model.load_state_dict(torch.load(CKPT, map_location=DEVICE)["model"])
torch.manual_seed(0)
prompt = torch.tensor([encode("ROMEO:")], device=DEVICE)
with autocast():
    out = generate(raw_model, prompt, max_new_tokens=500, temperature=0.8, top_k=20)
print(decode(out[0].tolist()))

为什么要"只在 val loss 创新低时保存"？ 1000 万参数的模型只有 100 万个字符可学，训练后期一定会开始背训练集（21.6 节会看到）。最终那一步的权重往往不是最好的，所以每次评估时比较 val loss，只保留最好的那一份。nanoGPT 的 shakespeare_char 配置也是这么做的。

21.5.3 GPU 加速四件套

设置	代码	作用
TF32	`allow_tf32 = True`	FP32 矩阵乘用 TF32 Tensor Core 计算，精度损失很小，速度数倍
bf16 混合精度	`torch.autocast(dtype=torch.bfloat16)`	前向/反向的大部分计算用 bf16；指数位与 FP32 相同，不需要 `GradScaler`
`torch.compile`	`model = torch.compile(raw_model)`	把模型编译成融合后的 Triton kernel，减少 kernel 数量和显存读写；首次编译约需几十秒
fused AdamW	`AdamW(..., fused=True)`	把所有参数的更新融合成少数几个 kernel，而不是每个参数一个

torch.compile 依赖 Triton，Windows 上默认关闭（os.name != "nt"）。

训练用编译后的 model，生成用 raw_model ：生成时序列长度每步都在变，编译后的模型遇到新 Shape 会触发重新编译，反而更慢。两者共享同一份参数，raw_model.state_dict() 保存的就是训练好的权重，不会带上 torch.compile 加的 _orig_mod. 前缀。

其他细节：

梯度裁剪 放在 backward() 之后、step() 之前：clip_grad_norm_ 把所有参数梯度的整体 L2 范数限制在 1.0 以内，偶发的大梯度不会把参数一次推飞。
beta_2=0.99（默认 0.999）：二阶矩估计对最近的梯度更敏感，小数据集、高学习率时更稳。
学习率手动设置 ：每步直接改 param_groups[i]["lr"]，比 LambdaLR 更直观，也方便打印。

21.6 运行结果

bash 复制代码

python train_gpt.py

text 复制代码

字符数 1,115,394，词表大小 65
device=cuda bf16=True compile=True 参数量 10.77M
初始 loss 理论值 ln(65) = 4.174
step     0 | train 4.3026 | val 4.2998 | gap -0.003 | lr 1.00e-05 | 2s  *saved
step   250 | train 1.9456 | val 2.0600 | gap 0.114 | lr 9.98e-04 | 6s  *saved
step   500 | train 1.5203 | val 1.7048 | gap 0.185 | lr 9.85e-04 | 11s  *saved
step   750 | train 1.3573 | val 1.5772 | gap 0.220 | lr 9.61e-04 | 15s  *saved
step  1000 | train 1.2652 | val 1.5118 | gap 0.247 | lr 9.27e-04 | 19s  *saved
step  1250 | train 1.2004 | val 1.4846 | gap 0.284 | lr 8.83e-04 | 24s  *saved
step  1500 | train 1.1417 | val 1.4796 | gap 0.338 | lr 8.31e-04 | 28s  *saved
step  1750 | train 1.0913 | val 1.4729 | gap 0.382 | lr 7.71e-04 | 32s  *saved
step  2000 | train 1.0405 | val 1.4766 | gap 0.436 | lr 7.05e-04 | 37s
step  2250 | train 0.9908 | val 1.4801 | gap 0.489 | lr 6.36e-04 | 41s
step  2500 | train 0.9449 | val 1.4998 | gap 0.555 | lr 5.64e-04 | 45s
step  2750 | train 0.8950 | val 1.5187 | gap 0.624 | lr 4.92e-04 | 49s
step  3000 | train 0.8444 | val 1.5473 | gap 0.703 | lr 4.22e-04 | 54s
step  3250 | train 0.7973 | val 1.5798 | gap 0.783 | lr 3.55e-04 | 58s
step  3500 | train 0.7540 | val 1.6051 | gap 0.851 | lr 2.93e-04 | 62s
step  3750 | train 0.7160 | val 1.6335 | gap 0.917 | lr 2.37e-04 | 66s
step  4000 | train 0.6813 | val 1.6578 | gap 0.976 | lr 1.89e-04 | 71s
step  4250 | train 0.6489 | val 1.6710 | gap 1.022 | lr 1.51e-04 | 75s
step  4500 | train 0.6244 | val 1.6832 | gap 1.059 | lr 1.23e-04 | 79s
step  4750 | train 0.6045 | val 1.7003 | gap 1.096 | lr 1.06e-04 | 84s
step  5000 | train 0.5906 | val 1.7255 | gap 1.135 | lr 1.00e-04 | 88s
峰值显存 1.61 GB
最佳 val loss 1.4729（step 1750），已保存到 gpt_shakespeare.pt

5000 步共 88 秒，峰值显存 1.61GB------这个规模的模型在 4090 上几乎不费力。

21.6.1 怎么读这些 loss？

基准线 ln⁡65=4.174\ln 65 = 4.174ln65=4.174 ：模型完全不懂、对 65 个字符均匀猜测时的交叉熵。step 0 的 4.30 已经很接近它，因为 N(0,0.02)\mathcal{N}(0, 0.02)N(0,0.02) 的小初始化让初始 logits 几乎都接近 0，输出分布接近均匀。
降到 2.06（step 250）：模型学会了字符频率和常见的字母组合，生成的内容"像英语字母"但不成词。
降到 1.47（step 1750） ：相当于每个位置平均在 e1.47≈4.4e^{1.47} ≈ 4.4e1.47≈4.4 个字符里犹豫，已经能写出大部分正确的单词、人名格式和换行结构。这和 nanoGPT 在同样配置下报告的最佳 val loss 1.4697 基本一致。

21.6.2 过拟合：gap 一路变大

#mermaid-svg-scjtUWx1CdXINGzG{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-scjtUWx1CdXINGzG .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-scjtUWx1CdXINGzG .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-scjtUWx1CdXINGzG .error-icon{fill:#552222;}#mermaid-svg-scjtUWx1CdXINGzG .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-scjtUWx1CdXINGzG .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-scjtUWx1CdXINGzG .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-scjtUWx1CdXINGzG .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-scjtUWx1CdXINGzG .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-scjtUWx1CdXINGzG .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-scjtUWx1CdXINGzG .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-scjtUWx1CdXINGzG .marker{fill:#333333;stroke:#333333;}#mermaid-svg-scjtUWx1CdXINGzG .marker.cross{stroke:#333333;}#mermaid-svg-scjtUWx1CdXINGzG svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-scjtUWx1CdXINGzG p{margin:0;}#mermaid-svg-scjtUWx1CdXINGzG .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-scjtUWx1CdXINGzG .cluster-label text{fill:#333;}#mermaid-svg-scjtUWx1CdXINGzG .cluster-label span{color:#333;}#mermaid-svg-scjtUWx1CdXINGzG .cluster-label span p{background-color:transparent;}#mermaid-svg-scjtUWx1CdXINGzG .label text,#mermaid-svg-scjtUWx1CdXINGzG span{fill:#333;color:#333;}#mermaid-svg-scjtUWx1CdXINGzG .node rect,#mermaid-svg-scjtUWx1CdXINGzG .node circle,#mermaid-svg-scjtUWx1CdXINGzG .node ellipse,#mermaid-svg-scjtUWx1CdXINGzG .node polygon,#mermaid-svg-scjtUWx1CdXINGzG .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-scjtUWx1CdXINGzG .rough-node .label text,#mermaid-svg-scjtUWx1CdXINGzG .node .label text,#mermaid-svg-scjtUWx1CdXINGzG .image-shape .label,#mermaid-svg-scjtUWx1CdXINGzG .icon-shape .label{text-anchor:middle;}#mermaid-svg-scjtUWx1CdXINGzG .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-scjtUWx1CdXINGzG .rough-node .label,#mermaid-svg-scjtUWx1CdXINGzG .node .label,#mermaid-svg-scjtUWx1CdXINGzG .image-shape .label,#mermaid-svg-scjtUWx1CdXINGzG .icon-shape .label{text-align:center;}#mermaid-svg-scjtUWx1CdXINGzG .node.clickable{cursor:pointer;}#mermaid-svg-scjtUWx1CdXINGzG .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-scjtUWx1CdXINGzG .arrowheadPath{fill:#333333;}#mermaid-svg-scjtUWx1CdXINGzG .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-scjtUWx1CdXINGzG .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-scjtUWx1CdXINGzG .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-scjtUWx1CdXINGzG .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-scjtUWx1CdXINGzG .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-scjtUWx1CdXINGzG .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-scjtUWx1CdXINGzG .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-scjtUWx1CdXINGzG .cluster text{fill:#333;}#mermaid-svg-scjtUWx1CdXINGzG .cluster span{color:#333;}#mermaid-svg-scjtUWx1CdXINGzG div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-scjtUWx1CdXINGzG .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-scjtUWx1CdXINGzG rect.text{fill:none;stroke-width:0;}#mermaid-svg-scjtUWx1CdXINGzG .icon-shape,#mermaid-svg-scjtUWx1CdXINGzG .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-scjtUWx1CdXINGzG .icon-shape p,#mermaid-svg-scjtUWx1CdXINGzG .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-scjtUWx1CdXINGzG .icon-shape .label rect,#mermaid-svg-scjtUWx1CdXINGzG .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-scjtUWx1CdXINGzG .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-scjtUWx1CdXINGzG .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-scjtUWx1CdXINGzG :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} step 0~1750

train ↓ val ↓

gap 0 → 0.38
step 1750

val 最低 1.473

保存 checkpoint
step 1750~5000

train ↓ val ↑

gap 0.38 → 1.14

阶段	train loss	val loss	gap（val − train）	模型在做什么
0 → 1750 步	4.30 → 1.09	4.30 → 1.47	0 → 0.38	学习可以泛化的规律：拼写、语法、剧本格式
1750 → 5000 步	1.09 → 0.59	1.47 → 1.73	0.38 → 1.14	背训练集：train loss 继续降，val loss 反而升

1000 万参数只配 100 万字符的数据，模型容量远超数据量，过拟合是必然的。Dropout 0.2 和 Weight Decay 0.1 能推迟它，但挡不住。应对方法：

早停（Early Stopping）：本章的做法，只保留 val loss 最低的 checkpoint；
更多数据（这是根本办法，也是大模型 Scaling Law 的出发点）；
更小的模型，或更强的正则化（练习 2）。

21.6.3 生成样本

用 step 1750 的最佳 checkpoint，temperature=0.8, top_k=20，以 ROMEO: 开头：

text 复制代码

ROMEO:
I am his ears and in his birth.
 
First Murderer:
No; I would were conferent for silkingly.
 
Second Murderer:
He did he speak to the rest of Caius Marcius
When he hath said to put the part bending of his lady.
 
First Murderer:
By the king, I can well send him for a parle.
 
Second Murderer:
He was a parlous cocked blood to send him;
and in him, is which he is subdue did by him?
 
Second Murderer:
Well, so for he delivers his good resolution:
'Zounds, let us go with him greater with the soul;
And,

一个字符级的小模型，90 秒训练，已经学会了：

剧本格式：角色名 + 冒号 + 换行 + 台词，台词之间空一行；
真实的角色和人名：First / Second Murderer（《理查三世》里的刺客）、Caius Marcius（《科利奥兰纳斯》的主角）；
莎士比亚式用词：'Zounds、parlous、hath。

但它不懂语义：句子是"像莎士比亚的胡话"。这正是第14章说的------Next Token Prediction 先学到的是形式，语义和知识需要大得多的模型和数据。

21.7 实验：Temperature 对生成的影响

python 复制代码

# ch21_gpt/common.py
"""加载 train_gpt.py 保存的 checkpoint"""
import torch
 
from gpt import GPT
 
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 
def load_model(path="gpt_shakespeare.pt"):
    ckpt = torch.load(path, map_location=DEVICE)
    model = GPT(**ckpt["config"]).to(DEVICE).eval()
    model.load_state_dict(ckpt["model"])
    chars = ckpt["chars"]
    stoi = {c: i for i, c in enumerate(chars)}
    encode = lambda s: [stoi[c] for c in s]
    decode = lambda ids: "".join(chars[int(i)] for i in ids)
    return model, encode, decode

python 复制代码

# ch21_gpt/temp_exp.py
"""Temperature 对生成的影响
 
运行：python temp_exp.py
"""
import torch
 
from common import DEVICE, load_model
from gpt import generate_with_cache
 
model, encode, decode = load_model()
prompt = torch.tensor([encode("KING RICHARD:\n")], device=DEVICE)
 
for t in (0, 0.5, 1.0, 1.5):
    torch.manual_seed(42)
    out = generate_with_cache(model, prompt, 150, temperature=t)
    print(f"===== temperature={t} =====")
    print(decode(out[0].tolist()))
    print()

（generate_with_cache 在 21.8 节实现，与 generate 结果等价。）

text 复制代码

===== temperature=0 =====
KING RICHARD:
What shall be the court? what shall I be so?
 
QUEEN ELIZABETH:
What shall we see the king hath so much a court?
 
KING RICHARD III:
What shall be the k
 
===== temperature=0.5 =====
KING RICHARD:
What shall be spoke with a month and where you shall be many with
him with his father and state him and stand up,
And he may be hanged by the market-s
 
===== temperature=1.0 =====
KING RICHARD:
What shall be youngest traitors are you,
Yourself have me of York lap with you. Proclaim,
The lusty courses of your true man, your has
Made a man term
 
===== temperature=1.5 =====
KING RICHARD:
I, till Richard's sworn, and distund he be:
Thereuoen, my gather, ladies Clifford, or and
The lusty, oby indeedined Plantagenet,
Wheps brotherleap He

Temperature	现象	原因
0（贪心）	每段台词都以 `What shall...` 开头，句式反复	每步都选最大概率，一旦进入高概率模式就很难跳出
0.5	通顺、保守，`him with his father` 这类高频搭配多	分布被压尖，低概率字符几乎选不到
1.0	最"像"原文，用词更丰富（`lusty courses`、`Proclaim`）	按模型学到的真实分布采样
1.5	拼写崩坏（`Thereuoen`、`indeedined`、`Wheps`）	分布被拉平，长尾的错误字符频繁被选中

贪心解码的重复问题 在大模型上同样存在，只是形式不同：模型越好，越不会出现 the the the 这样的逐词循环，而是变成句式、段落层面的重复。这也是为什么对话类应用通常把 temperature 设在 0.7 左右，并配合 Top-p 与重复惩罚（Repetition Penalty）。

21.8 实现 KV Cache

21.8.1 朴素生成浪费在哪？

generate 每生成一个字符，都把整个上下文重新跑一遍。生成第 t 个字符时，前 t−1 个位置的 K、V 在上一步已经算过一次，这一步又原样重算------而由于 Causal Mask，前面位置的输出根本不受新 Token 影响（第15章 15.5.3）。
#mermaid-svg-2axFBLMd62Dwz1wE{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-2axFBLMd62Dwz1wE .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-2axFBLMd62Dwz1wE .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-2axFBLMd62Dwz1wE .error-icon{fill:#552222;}#mermaid-svg-2axFBLMd62Dwz1wE .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-2axFBLMd62Dwz1wE .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-2axFBLMd62Dwz1wE .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-2axFBLMd62Dwz1wE .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-2axFBLMd62Dwz1wE .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-2axFBLMd62Dwz1wE .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-2axFBLMd62Dwz1wE .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-2axFBLMd62Dwz1wE .marker{fill:#333333;stroke:#333333;}#mermaid-svg-2axFBLMd62Dwz1wE .marker.cross{stroke:#333333;}#mermaid-svg-2axFBLMd62Dwz1wE svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-2axFBLMd62Dwz1wE p{margin:0;}#mermaid-svg-2axFBLMd62Dwz1wE .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-2axFBLMd62Dwz1wE .cluster-label text{fill:#333;}#mermaid-svg-2axFBLMd62Dwz1wE .cluster-label span{color:#333;}#mermaid-svg-2axFBLMd62Dwz1wE .cluster-label span p{background-color:transparent;}#mermaid-svg-2axFBLMd62Dwz1wE .label text,#mermaid-svg-2axFBLMd62Dwz1wE span{fill:#333;color:#333;}#mermaid-svg-2axFBLMd62Dwz1wE .node rect,#mermaid-svg-2axFBLMd62Dwz1wE .node circle,#mermaid-svg-2axFBLMd62Dwz1wE .node ellipse,#mermaid-svg-2axFBLMd62Dwz1wE .node polygon,#mermaid-svg-2axFBLMd62Dwz1wE .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-2axFBLMd62Dwz1wE .rough-node .label text,#mermaid-svg-2axFBLMd62Dwz1wE .node .label text,#mermaid-svg-2axFBLMd62Dwz1wE .image-shape .label,#mermaid-svg-2axFBLMd62Dwz1wE .icon-shape .label{text-anchor:middle;}#mermaid-svg-2axFBLMd62Dwz1wE .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-2axFBLMd62Dwz1wE .rough-node .label,#mermaid-svg-2axFBLMd62Dwz1wE .node .label,#mermaid-svg-2axFBLMd62Dwz1wE .image-shape .label,#mermaid-svg-2axFBLMd62Dwz1wE .icon-shape .label{text-align:center;}#mermaid-svg-2axFBLMd62Dwz1wE .node.clickable{cursor:pointer;}#mermaid-svg-2axFBLMd62Dwz1wE .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-2axFBLMd62Dwz1wE .arrowheadPath{fill:#333333;}#mermaid-svg-2axFBLMd62Dwz1wE .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-2axFBLMd62Dwz1wE .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-2axFBLMd62Dwz1wE .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-2axFBLMd62Dwz1wE .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-2axFBLMd62Dwz1wE .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-2axFBLMd62Dwz1wE .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-2axFBLMd62Dwz1wE .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-2axFBLMd62Dwz1wE .cluster text{fill:#333;}#mermaid-svg-2axFBLMd62Dwz1wE .cluster span{color:#333;}#mermaid-svg-2axFBLMd62Dwz1wE div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-2axFBLMd62Dwz1wE .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-2axFBLMd62Dwz1wE rect.text{fill:none;stroke-width:0;}#mermaid-svg-2axFBLMd62Dwz1wE .icon-shape,#mermaid-svg-2axFBLMd62Dwz1wE .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-2axFBLMd62Dwz1wE .icon-shape p,#mermaid-svg-2axFBLMd62Dwz1wE .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-2axFBLMd62Dwz1wE .icon-shape .label rect,#mermaid-svg-2axFBLMd62Dwz1wE .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-2axFBLMd62Dwz1wE .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-2axFBLMd62Dwz1wE .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-2axFBLMd62Dwz1wE :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} KV Cache
Prefill：算 P 个位置

缓存每层 K/V
Decode：只算 1 个新位置

K/V 追加进缓存
Decode：只算 1 个新位置
朴素生成
第1步：算 P 个位置
第2步：算 P+1 个位置
第3步：算 P+2 个位置

21.8.2 代码

模型部分在 21.3 节已经预留了接口，生成函数只需要管理 Cache：

python 复制代码

# ch21_gpt/gpt.py（续）
@torch.no_grad()
def generate_with_cache(model, idx, max_new_tokens, temperature=1.0, top_k=None):
    """KV Cache 生成：Prompt 跑一次（Prefill），之后每步只算 1 个新 Token（Decode）"""
    model.eval()
    total = idx.size(1) + max_new_tokens
    assert total <= model.block_size, "绝对位置编码 + Cache 不能超过 block_size"
 
    caches = [{"k": None, "v": None} for _ in model.blocks]
    logits, caches = model(idx, caches=caches, start_pos=0)            # Prefill
    out = idx
    for i in range(max_new_tokens):
        nxt = sample_next(logits[:, -1], temperature, top_k)
        out = torch.cat([out, nxt], dim=1)
        if i < max_new_tokens - 1:                                     # Decode：只喂新 Token
            logits, caches = model(nxt, caches=caches, start_pos=out.size(1) - 1)
    return out

每层 Cache 的 Shape 变化：

阶段	输入	本层新算的 K	Cache 中的 K
Prefill（Prompt 长 P）	$B, P$	$B, h, P, d_k$	$B, h, P, d_k$
Decode 第 1 步	$B, 1$	$B, h, 1, d_k$	$B, h, P+1, d_k$
Decode 第 n 步	$B, 1$	$B, h, 1, d_k$	$B, h, P+n, d_k$

为什么有 assert total <= block_size？ 朴素生成超过窗口时可以简单截掉最前面的字符重新计算；但 Cache 里的 K 已经带上了"第 0、1、2... 个位置"的位置编码，截掉开头后整体位置都错了，不能直接复用。这是绝对位置编码的局限之一，RoPE（第15章）配合 Sliding Window 等技术才能更好地处理。

21.8.3 验证：结果必须逐字一致

KV Cache 是纯工程优化，不应该改变任何计算结果：

python 复制代码

# ch21_gpt/test_cache.py
"""KV Cache：① 与朴素生成结果逐字一致  ② 速度对比
 
运行：python test_cache.py
"""
import time
 
import torch
 
from common import DEVICE, load_model
from gpt import generate, generate_with_cache
 
# 一致性测试用纯 FP32，关闭 TF32：两条路径的矩阵形状不同，低精度下舍入误差可能让 argmax 翻转
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False
 
model, encode, decode = load_model()
prompt = torch.tensor([encode("ROMEO:")], device=DEVICE)
N = model.block_size - prompt.size(1)                    # 一直生成到上下文窗口填满
 
a = generate(model, prompt, N, temperature=0)            # 贪心，结果确定
b = generate_with_cache(model, prompt, N, temperature=0)
same = torch.equal(a, b)
print(f"生成 {N} 个 Token，两种方式输出一致: {same}")
if not same:
    first = (a != b).nonzero()[0, 1].item()
    print(f"  第一次不一致出现在位置 {first}（若只在很靠后才分叉，通常是浮点舍入，不是 bug）")
 
 
# 分块 Prefill：先喂前 40 个 Token，再带着 Cache 一次喂后 24 个（走 gpt.py 里"显式 Mask"分支）
with torch.no_grad():
    seq = a[:, :64]
    full, _ = model(seq)
    caches = [{"k": None, "v": None} for _ in model.blocks]
    _, caches = model(seq[:, :40], caches=caches, start_pos=0)
    part, _ = model(seq[:, 40:], caches=caches, start_pos=40)
print("分块 Prefill 与整段前向一致:", torch.allclose(full[:, 40:], part, atol=1e-4))
 
 
def bench(fn, repeat=3):
    fn(model, prompt, N, temperature=0)                  # 预热：CUDA 初始化、内核选择
    if DEVICE == "cuda":
        torch.cuda.synchronize()                         # GPU 是异步执行的，计时前后都要同步
    t = time.perf_counter()
    for _ in range(repeat):
        fn(model, prompt, N, temperature=0)
    if DEVICE == "cuda":
        torch.cuda.synchronize()
    return (time.perf_counter() - t) / repeat
 
 
for precision in ("fp32", "bf16"):
    enabled = precision == "bf16" and DEVICE == "cuda"
    if precision == "bf16" and not enabled:
        continue
    with torch.autocast(device_type=DEVICE, dtype=torch.bfloat16, enabled=enabled):
        t1, t2 = bench(generate), bench(generate_with_cache)
    print(f"[{precision}] generate {t1:.2f}s | generate_with_cache {t2:.2f}s | "
          f"加速 {t1 / t2:.1f}x | {N / t2:.0f} tokens/s")

text 复制代码

生成 250 个 Token，两种方式输出一致: True
分块 Prefill 与整段前向一致: True
[fp32] generate 0.14s | generate_with_cache 0.15s | 加速 0.9x | 1644 tokens/s
[bf16] generate 0.16s | generate_with_cache 0.16s | 加速 1.0x | 1527 tokens/s

逐字一致：贪心生成 250 个字符（填满整个上下文窗口），两种方式完全相同。测试时关掉了 TF32：两条路径的矩阵形状不同，低精度下的舍入误差可能让某个 argmax 翻转，导致后面全部分叉------那不是 bug，但会干扰判断。
分块 Prefill 一致 ：先喂 40 个 Token 建 Cache，再一次喂 24 个，走的是 21.3.1 节"显式右下角 Mask"那个分支，结果与整段 64 个 Token 一起前向相同。如果这里误用 is_causal=True，这个测试会失败。
但是一点都没快，甚至 fp32 下还慢了一点。

21.8.4 为什么没有加速？

计算量的差距其实很大：Prompt 6 个字符、生成 250 个，朴素方式累计要处理约 ∑t=6255t≈3.3\sum_{t=6}^{255} t ≈ 3.3∑t=6255t≈3.3 万个位置，KV Cache 只处理 6+249=2556 + 249 = 2556+249=255 个，线性层计算量少了约 130 倍。时间却一样，原因是：

在 4090 上跑一个 1000 万参数、上下文 256、batch=1 的模型，GPU 每一步的计算只需要几微秒，绝大部分时间在等 CPU 把下一个 kernel 发过来（launch-bound）。

每生成一个 Token，两种方式都要依次发射几十个 kernel（6 层 × 每层若干个矩阵乘、LayerNorm、Attention、加法），这部分开销两者一样多；省下的那点计算量相对 kernel 发射开销可以忽略。

所以要看到 KV Cache 的效果，必须让计算量本身成为瓶颈。

21.9 实验：KV Cache 什么时候才有用？

用一个随机初始化的 GPT-2 Small 尺寸模型（12 层、768 维，只测速度，不需要训练），扫描 batch 大小和生成长度：

python 复制代码

# ch21_gpt/bench_cache.py
"""KV Cache 在什么条件下才显出优势？
 
test_cache.py 里 10M 参数、250 个 Token、batch=1 的小模型两种方式一样快：
GPU 每一步都在等 CPU 发射 kernel（launch-bound），计算量根本不是瓶颈。
这里用随机初始化的 GPT-2 Small 尺寸模型（只测速度，不需要训练），扫描 batch 和生成长度。
 
运行：python bench_cache.py
"""
import time
 
import torch
 
from gpt import GPT, generate, generate_with_cache
 
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
torch.backends.cuda.matmul.allow_tf32 = True
 
# GPT-2 Small：12 层、768 维、12 头、上下文 1024
model = GPT(vocab_size=65, block_size=1024, d_model=768, num_heads=12, num_layers=12,
            dropout=0.0).to(DEVICE).eval()
print(f"device={DEVICE}  参数量 {sum(p.numel() for p in model.parameters()) / 1e6:.1f}M")
 
 
def timed(fn, prompt, n):
    if DEVICE == "cuda":
        torch.cuda.synchronize()
        torch.cuda.reset_peak_memory_stats()
    t = time.perf_counter()
    with torch.autocast(device_type=DEVICE, dtype=torch.bfloat16, enabled=DEVICE == "cuda"):
        fn(model, prompt, n, temperature=0)
    if DEVICE == "cuda":
        torch.cuda.synchronize()
    mem = torch.cuda.max_memory_allocated() / 2**30 if DEVICE == "cuda" else 0.0
    return time.perf_counter() - t, mem
 
 
prompt_len = 16
timed(generate_with_cache, torch.zeros(1, prompt_len, dtype=torch.long, device=DEVICE), 8)   # 预热
timed(generate, torch.zeros(1, prompt_len, dtype=torch.long, device=DEVICE), 8)
 
print(f"{'batch':>5} {'生成长度':>8} | {'朴素':>8} {'KV Cache':>9} {'加速':>6} | "
      f"{'朴素显存':>8} {'Cache显存':>9} | {'Cache tokens/s':>14}")
for batch in (1, 8, 32):
    for n in (256, 1000):
        prompt = torch.randint(0, 65, (batch, prompt_len), device=DEVICE)
        t_naive, m_naive = timed(generate, prompt, n)
        t_cache, m_cache = timed(generate_with_cache, prompt, n)
        print(f"{batch:>5} {n:>8} | {t_naive:>7.2f}s {t_cache:>8.2f}s {t_naive / t_cache:>5.1f}x | "
              f"{m_naive:>7.2f}G {m_cache:>8.2f}G | {batch * n / t_cache:>14.0f}")

实测（RTX 4090，bf16）：

text 复制代码

device=cuda  参数量 85.9M
batch     生成长度 |       朴素  KV Cache     加速 |     朴素显存   Cache显存 | Cache tokens/s
    1      256 |    0.31s     0.32s   1.0x |    0.51G     0.52G |            809
    1     1000 |    1.49s     1.24s   1.2x |    0.52G     0.58G |            805
    8      256 |    0.56s     0.32s   1.7x |    0.55G     0.65G |           6330
    8     1000 |    7.09s     1.43s   4.9x |    0.68G     1.06G |           5580
   32      256 |    1.93s     0.39s   4.9x |    0.69G     1.10G |          20766
   32     1000 |   29.05s     3.20s   9.1x |    1.20G     2.73G |           9998

21.9.1 加速比：计算量越大，KV Cache 越有用

条件	加速	解释
batch=1，生成 256	1.0x	依然 launch-bound，两者都在等 kernel 发射
batch=1，生成 1000	1.2x	序列变长，朴素方式的计算量开始冒头
batch=8，生成 1000	4.9x	朴素方式的计算量按 B×T2B \times T^2B×T2 增长，成为瓶颈
batch=32，生成 1000	9.1x	朴素 29 秒 vs Cache 3.2 秒

注意 KV Cache 那一列：batch 从 1 增加到 32，耗时只从 0.32s 增加到 0.39s（生成 256 时）。这说明 Decode 阶段 GPU 远没有吃满，多个请求合成一个 batch 一起解码几乎是"免费"的------这正是 vLLM 等推理框架做 Continuous Batching 的原因。

21.9.2 显存：用空间换时间

batch=32、生成 1000 时，Cache 版峰值显存 2.73GB，朴素版 1.20GB。KV Cache 本身的大小可以直接算出来：

12⏟层×2⏟K,V×32⏟batch×1016⏟长度×768⏟dmodel×2 B⏟bf16≈1.2 GB \underbrace{12}{\text{层}} \times \underbrace{2}{K,V} \times \underbrace{32}{batch} \times \underbrace{1016}{\text{长度}} \times \underbrace{768}{d{model}} \times \underbrace{2\text{ B}}_{bf16} ≈ 1.2\text{ GB} 层 12×K,V 2×batch 32×长度 1016×dmodel 768×bf16 2 B≈1.2 GB

多出来的部分基本就是 Cache，再加上 torch.cat 每一步分配新张量时的临时副本。对 7B 模型、4K 上下文，这个数字会到每个请求 GB 级别------这就是第15章 GQA 要把 K/V 头数砍掉的原因。

21.9.3 吞吐：生成越长，越慢

同样 batch=32，生成 256 时吞吐 20766 tokens/s，生成 1000 时只有 9998 tokens/s。两个原因：

每步 Attention 要扫描的历史越来越长，这部分是 O(T)O(T)O(T) 的；
本章的实现每一步都 torch.cat 出一个新的 Cache 张量，要把整个旧 Cache 复制一遍，累计显存搬运量是 O(T2)O(T^2)O(T2)。

工业实现会预分配固定长度的 Cache，每步只写入新位置（练习 6）；vLLM 的 PagedAttention 进一步把 Cache 切成固定大小的"页"，解决多请求共享显存时的碎片问题。

21.9.4 CUDA Graph：彻底消除 launch-bound

21.8.4 节说过，batch=1 时 KV Cache 没有加速，原因是 GPU 每一步都在等 CPU 发射 kernel。CUDA Graph 是解决这个问题的根本手段。

原理：

普通执行模式下，每生成一个 Token，CPU 要依次向 GPU 发射几十个 kernel（矩阵乘、LayerNorm、Attention......），每次发射都有固定的 CPU 开销（约几微秒）。当计算本身只需要几微秒时，发射开销就成了瓶颈：

复制代码

普通模式：
CPU: [发射k1]--[发射k2]--[发射k3]-- ...
GPU:      [k1]      [k2]      [k3]  ...
          ↑等待CPU发射，GPU大量空闲
 
CUDA Graph 模式：
CPU: [录制]→[一次提交整张图]
GPU: [k1][k2][k3][k4]...  ← 连续执行，无等待

CUDA Graph 把一整步 Decode 的所有 kernel 录制成一张"图"，之后每步只需 CPU 发射一次"执行这张图"的指令，GPU 就能连续跑完所有 kernel，彻底消除发射开销。

使用条件：

CUDA Graph 要求每次执行时输入输出的形状和显存地址完全固定 ，否则需要重新录制。这正是 Decode 阶段（每步输入都是 [B, 1]）天然满足、而 Prefill 阶段（输入长度变化）不满足的原因：

阶段	输入形状	适合 CUDA Graph？
Prefill	`[B, P]`，P 每次不同	不适合
Decode	`[B, 1]`，固定	适合

因此工业实现通常只对 Decode 步录制 CUDA Graph，Prefill 仍走普通模式。

PyTorch 的使用方式：

torch.compile(mode="reduce-overhead") 内部会自动尝试使用 CUDA Graph，是最简单的启用方式：

python 复制代码

# 预分配固定形状的 Cache（形状不变才能用 CUDA Graph）
decode_step = torch.compile(model, mode="reduce-overhead")
 
# Prefill：普通模式，形状不固定
logits, caches = model(prompt, caches=caches, start_pos=0)
 
# Decode：编译后的模型，形状固定为 [B, 1]，触发 CUDA Graph
for i in range(max_new_tokens):
    logits, caches = decode_step(nxt, caches=caches, start_pos=prompt_len + i)

首次执行时 PyTorch 会自动录制，后续每步直接回放，batch=1 时也能接近 GPU 算力上限。具体实现见练习 6。

21.10 Hugging Face 源码定位

本章组件	Hugging Face 对应位置（`src/transformers/`）	说明
`CausalSelfAttention`	`models/gpt2/modeling_gpt2.py` → `GPT2Attention`	同样用一个 `c_attn` 投影出 Q/K/V；SDPA 后端在 `integrations/sdpa_attention.py`
`GPTBlock`	`models/gpt2/modeling_gpt2.py` → `GPT2Block`	Pre-Norm：`ln_1` → attn → `ln_2` → mlp
`GPT`	`models/gpt2/modeling_gpt2.py` → `GPT2LMHeadModel`	`lm_head` 与 `wte` 共享权重
KV Cache	`cache_utils.py` → `DynamicCache` / `StaticCache`	`StaticCache` 就是预分配版本，可配合 `torch.compile`
`generate` / 采样	`generation/utils.py` → `GenerationMixin.generate`	Temperature/Top-k 在 `generation/logits_process.py`

nanoGPT（github.com/karpathy/nanoGPT）的 model.py 只有 300 多行，和本章代码几乎一一对应，非常适合作为下一步阅读材料。

21.11 常见坑

现象	原因	解决
带 Cache 一次喂多个 Token 后结果错乱	SDPA 的 `is_causal=True` 在 Tq≠TkT_q \ne T_kTq=Tk 时按左上角对齐	显式构造右下角对齐的 Mask（21.3.1）
换成 SDPA 后模型学不会	SDPA 的 bool mask 是 True = 可见，传反了	自回归直接用 `is_causal=True`
训练很久后生成质量反而变差	过拟合，保存的是最后一步的权重	按 val loss 保存最佳 checkpoint（21.5.2）
加载 checkpoint 报 `Unexpected key(s): _orig_mod.xxx`	保存了 `torch.compile` 后模型的 `state_dict`	保存 `raw_model.state_dict()`
生成时越来越慢、或频繁打印重新编译警告	用编译后的模型做变长生成	生成用未编译的 `raw_model`
KV Cache 生成结果与朴素生成不一致	`start_pos` 没传，新 Token 用了位置 0 的位置编码	传入新 Token 在整段序列中的真实位置
一致性测试偶尔在很靠后的位置分叉	bf16 / TF32 舍入误差让某个 argmax 翻转	测试时用纯 FP32 并关闭 TF32
用 KV Cache 却没有加速	小模型 + batch=1，GPU 在等 kernel 发射	正常现象；增大 batch / 长度，或用 CUDA Graph（21.9.4）
生成到一定长度后报 `index out of range`	位置编号超过了 `pos_emb` 的 `block_size`	朴素生成截断上下文；Cache 生成限制总长
计时结果忽大忽小	GPU 异步执行，没等 kernel 跑完就停表	计时前后 `torch.cuda.synchronize()`

本章练习与实验

调参对抗过拟合 ：把 DROPOUT 改成 0.3、或 N_LAYER 改成 4，最佳 val loss 和出现的步数如何变化？
更多数据 ：把多份公开领域的英文文本拼进 input.txt（注意词表会变），过拟合是否推迟？
Top-p 采样 ：在 sample_next 中实现 Nucleus Sampling（第14章 14.12）。
换成 RMSNorm + SwiGLU + RoPE ：参照第9、8、15章，把 GPT Block 改成 Llama 风格，对比 loss 曲线。用 RoPE 之后，21.8.2 节的 assert 还必要吗？
BPE Tokenizer ：用 tiktoken 的 GPT-2 编码替换字符级 Tokenizer。词表变成 50257 后，参数量主要花在哪里了？
预分配 Cache ：把 Cache 改成长度固定为 block_size 的张量，每步用切片赋值写入新位置，再用 torch.compile(mode="reduce-overhead")（内部使用 CUDA Graph）编译 Decode 步，重跑 bench_cache.py，batch=1 时能快多少？
画出过拟合曲线：把 21.6 节的日志解析出来，用 matplotlib 画 train / val loss 随 step 的变化，标出最佳点。

本章总结

#mermaid-svg-hUkqAxA9OLdu2Wup{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-hUkqAxA9OLdu2Wup .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-hUkqAxA9OLdu2Wup .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-hUkqAxA9OLdu2Wup .error-icon{fill:#552222;}#mermaid-svg-hUkqAxA9OLdu2Wup .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-hUkqAxA9OLdu2Wup .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-hUkqAxA9OLdu2Wup .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-hUkqAxA9OLdu2Wup .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-hUkqAxA9OLdu2Wup .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-hUkqAxA9OLdu2Wup .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-hUkqAxA9OLdu2Wup .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-hUkqAxA9OLdu2Wup .marker{fill:#333333;stroke:#333333;}#mermaid-svg-hUkqAxA9OLdu2Wup .marker.cross{stroke:#333333;}#mermaid-svg-hUkqAxA9OLdu2Wup svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-hUkqAxA9OLdu2Wup p{margin:0;}#mermaid-svg-hUkqAxA9OLdu2Wup .edge{stroke-width:3;}#mermaid-svg-hUkqAxA9OLdu2Wup .section--1 rect,#mermaid-svg-hUkqAxA9OLdu2Wup .section--1 path,#mermaid-svg-hUkqAxA9OLdu2Wup .section--1 circle,#mermaid-svg-hUkqAxA9OLdu2Wup .section--1 polygon,#mermaid-svg-hUkqAxA9OLdu2Wup .section--1 path{fill:hsl(240, 100%, 76.2745098039%);}#mermaid-svg-hUkqAxA9OLdu2Wup .section--1 text{fill:#ffffff;}#mermaid-svg-hUkqAxA9OLdu2Wup .node-icon--1{font-size:40px;color:#ffffff;}#mermaid-svg-hUkqAxA9OLdu2Wup .section-edge--1{stroke:hsl(240, 100%, 76.2745098039%);}#mermaid-svg-hUkqAxA9OLdu2Wup .edge-depth--1{stroke-width:17;}#mermaid-svg-hUkqAxA9OLdu2Wup .section--1 line{stroke:hsl(60, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-hUkqAxA9OLdu2Wup .disabled,#mermaid-svg-hUkqAxA9OLdu2Wup .disabled circle,#mermaid-svg-hUkqAxA9OLdu2Wup .disabled text{fill:lightgray;}#mermaid-svg-hUkqAxA9OLdu2Wup .disabled text{fill:#efefef;}#mermaid-svg-hUkqAxA9OLdu2Wup .section-0 rect,#mermaid-svg-hUkqAxA9OLdu2Wup .section-0 path,#mermaid-svg-hUkqAxA9OLdu2Wup .section-0 circle,#mermaid-svg-hUkqAxA9OLdu2Wup .section-0 polygon,#mermaid-svg-hUkqAxA9OLdu2Wup .section-0 path{fill:hsl(60, 100%, 73.5294117647%);}#mermaid-svg-hUkqAxA9OLdu2Wup .section-0 text{fill:black;}#mermaid-svg-hUkqAxA9OLdu2Wup .node-icon-0{font-size:40px;color:black;}#mermaid-svg-hUkqAxA9OLdu2Wup .section-edge-0{stroke:hsl(60, 100%, 73.5294117647%);}#mermaid-svg-hUkqAxA9OLdu2Wup .edge-depth-0{stroke-width:14;}#mermaid-svg-hUkqAxA9OLdu2Wup .section-0 line{stroke:hsl(240, 100%, 83.5294117647%);stroke-width:3;}#mermaid-svg-hUkqAxA9OLdu2Wup .disabled,#mermaid-svg-hUkqAxA9OLdu2Wup .disabled circle,#mermaid-svg-hUkqAxA9OLdu2Wup .disabled text{fill:lightgray;}#mermaid-svg-hUkqAxA9OLdu2Wup .disabled text{fill:#efefef;}#mermaid-svg-hUkqAxA9OLdu2Wup .section-1 rect,#mermaid-svg-hUkqAxA9OLdu2Wup .section-1 path,#mermaid-svg-hUkqAxA9OLdu2Wup .section-1 circle,#mermaid-svg-hUkqAxA9OLdu2Wup .section-1 polygon,#mermaid-svg-hUkqAxA9OLdu2Wup .section-1 path{fill:hsl(80, 100%, 76.2745098039%);}#mermaid-svg-hUkqAxA9OLdu2Wup .section-1 text{fill:black;}#mermaid-svg-hUkqAxA9OLdu2Wup .node-icon-1{font-size:40px;color:black;}#mermaid-svg-hUkqAxA9OLdu2Wup .section-edge-1{stroke:hsl(80, 100%, 76.2745098039%);}#mermaid-svg-hUkqAxA9OLdu2Wup .edge-depth-1{stroke-width:11;}#mermaid-svg-hUkqAxA9OLdu2Wup .section-1 line{stroke:hsl(260, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-hUkqAxA9OLdu2Wup .disabled,#mermaid-svg-hUkqAxA9OLdu2Wup .disabled circle,#mermaid-svg-hUkqAxA9OLdu2Wup .disabled text{fill:lightgray;}#mermaid-svg-hUkqAxA9OLdu2Wup .disabled text{fill:#efefef;}#mermaid-svg-hUkqAxA9OLdu2Wup .section-2 rect,#mermaid-svg-hUkqAxA9OLdu2Wup .section-2 path,#mermaid-svg-hUkqAxA9OLdu2Wup .section-2 circle,#mermaid-svg-hUkqAxA9OLdu2Wup .section-2 polygon,#mermaid-svg-hUkqAxA9OLdu2Wup .section-2 path{fill:hsl(270, 100%, 76.2745098039%);}#mermaid-svg-hUkqAxA9OLdu2Wup .section-2 text{fill:#ffffff;}#mermaid-svg-hUkqAxA9OLdu2Wup .node-icon-2{font-size:40px;color:#ffffff;}#mermaid-svg-hUkqAxA9OLdu2Wup .section-edge-2{stroke:hsl(270, 100%, 76.2745098039%);}#mermaid-svg-hUkqAxA9OLdu2Wup .edge-depth-2{stroke-width:8;}#mermaid-svg-hUkqAxA9OLdu2Wup .section-2 line{stroke:hsl(90, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-hUkqAxA9OLdu2Wup .disabled,#mermaid-svg-hUkqAxA9OLdu2Wup .disabled circle,#mermaid-svg-hUkqAxA9OLdu2Wup .disabled text{fill:lightgray;}#mermaid-svg-hUkqAxA9OLdu2Wup .disabled text{fill:#efefef;}#mermaid-svg-hUkqAxA9OLdu2Wup .section-3 rect,#mermaid-svg-hUkqAxA9OLdu2Wup .section-3 path,#mermaid-svg-hUkqAxA9OLdu2Wup .section-3 circle,#mermaid-svg-hUkqAxA9OLdu2Wup .section-3 polygon,#mermaid-svg-hUkqAxA9OLdu2Wup .section-3 path{fill:hsl(300, 100%, 76.2745098039%);}#mermaid-svg-hUkqAxA9OLdu2Wup .section-3 text{fill:black;}#mermaid-svg-hUkqAxA9OLdu2Wup .node-icon-3{font-size:40px;color:black;}#mermaid-svg-hUkqAxA9OLdu2Wup .section-edge-3{stroke:hsl(300, 100%, 76.2745098039%);}#mermaid-svg-hUkqAxA9OLdu2Wup .edge-depth-3{stroke-width:5;}#mermaid-svg-hUkqAxA9OLdu2Wup .section-3 line{stroke:hsl(120, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-hUkqAxA9OLdu2Wup .disabled,#mermaid-svg-hUkqAxA9OLdu2Wup .disabled circle,#mermaid-svg-hUkqAxA9OLdu2Wup .disabled text{fill:lightgray;}#mermaid-svg-hUkqAxA9OLdu2Wup .disabled text{fill:#efefef;}#mermaid-svg-hUkqAxA9OLdu2Wup .section-4 rect,#mermaid-svg-hUkqAxA9OLdu2Wup .section-4 path,#mermaid-svg-hUkqAxA9OLdu2Wup .section-4 circle,#mermaid-svg-hUkqAxA9OLdu2Wup .section-4 polygon,#mermaid-svg-hUkqAxA9OLdu2Wup .section-4 path{fill:hsl(330, 100%, 76.2745098039%);}#mermaid-svg-hUkqAxA9OLdu2Wup .section-4 text{fill:black;}#mermaid-svg-hUkqAxA9OLdu2Wup .node-icon-4{font-size:40px;color:black;}#mermaid-svg-hUkqAxA9OLdu2Wup .section-edge-4{stroke:hsl(330, 100%, 76.2745098039%);}#mermaid-svg-hUkqAxA9OLdu2Wup .edge-depth-4{stroke-width:2;}#mermaid-svg-hUkqAxA9OLdu2Wup .section-4 line{stroke:hsl(150, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-hUkqAxA9OLdu2Wup .disabled,#mermaid-svg-hUkqAxA9OLdu2Wup .disabled circle,#mermaid-svg-hUkqAxA9OLdu2Wup .disabled text{fill:lightgray;}#mermaid-svg-hUkqAxA9OLdu2Wup .disabled text{fill:#efefef;}#mermaid-svg-hUkqAxA9OLdu2Wup .section-5 rect,#mermaid-svg-hUkqAxA9OLdu2Wup .section-5 path,#mermaid-svg-hUkqAxA9OLdu2Wup .section-5 circle,#mermaid-svg-hUkqAxA9OLdu2Wup .section-5 polygon,#mermaid-svg-hUkqAxA9OLdu2Wup .section-5 path{fill:hsl(0, 100%, 76.2745098039%);}#mermaid-svg-hUkqAxA9OLdu2Wup .section-5 text{fill:black;}#mermaid-svg-hUkqAxA9OLdu2Wup .node-icon-5{font-size:40px;color:black;}#mermaid-svg-hUkqAxA9OLdu2Wup .section-edge-5{stroke:hsl(0, 100%, 76.2745098039%);}#mermaid-svg-hUkqAxA9OLdu2Wup .edge-depth-5{stroke-width:-1;}#mermaid-svg-hUkqAxA9OLdu2Wup .section-5 line{stroke:hsl(180, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-hUkqAxA9OLdu2Wup .disabled,#mermaid-svg-hUkqAxA9OLdu2Wup .disabled circle,#mermaid-svg-hUkqAxA9OLdu2Wup .disabled text{fill:lightgray;}#mermaid-svg-hUkqAxA9OLdu2Wup .disabled text{fill:#efefef;}#mermaid-svg-hUkqAxA9OLdu2Wup .section-6 rect,#mermaid-svg-hUkqAxA9OLdu2Wup .section-6 path,#mermaid-svg-hUkqAxA9OLdu2Wup .section-6 circle,#mermaid-svg-hUkqAxA9OLdu2Wup .section-6 polygon,#mermaid-svg-hUkqAxA9OLdu2Wup .section-6 path{fill:hsl(30, 100%, 76.2745098039%);}#mermaid-svg-hUkqAxA9OLdu2Wup .section-6 text{fill:black;}#mermaid-svg-hUkqAxA9OLdu2Wup .node-icon-6{font-size:40px;color:black;}#mermaid-svg-hUkqAxA9OLdu2Wup .section-edge-6{stroke:hsl(30, 100%, 76.2745098039%);}#mermaid-svg-hUkqAxA9OLdu2Wup .edge-depth-6{stroke-width:-4;}#mermaid-svg-hUkqAxA9OLdu2Wup .section-6 line{stroke:hsl(210, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-hUkqAxA9OLdu2Wup .disabled,#mermaid-svg-hUkqAxA9OLdu2Wup .disabled circle,#mermaid-svg-hUkqAxA9OLdu2Wup .disabled text{fill:lightgray;}#mermaid-svg-hUkqAxA9OLdu2Wup .disabled text{fill:#efefef;}#mermaid-svg-hUkqAxA9OLdu2Wup .section-7 rect,#mermaid-svg-hUkqAxA9OLdu2Wup .section-7 path,#mermaid-svg-hUkqAxA9OLdu2Wup .section-7 circle,#mermaid-svg-hUkqAxA9OLdu2Wup .section-7 polygon,#mermaid-svg-hUkqAxA9OLdu2Wup .section-7 path{fill:hsl(90, 100%, 76.2745098039%);}#mermaid-svg-hUkqAxA9OLdu2Wup .section-7 text{fill:black;}#mermaid-svg-hUkqAxA9OLdu2Wup .node-icon-7{font-size:40px;color:black;}#mermaid-svg-hUkqAxA9OLdu2Wup .section-edge-7{stroke:hsl(90, 100%, 76.2745098039%);}#mermaid-svg-hUkqAxA9OLdu2Wup .edge-depth-7{stroke-width:-7;}#mermaid-svg-hUkqAxA9OLdu2Wup .section-7 line{stroke:hsl(270, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-hUkqAxA9OLdu2Wup .disabled,#mermaid-svg-hUkqAxA9OLdu2Wup .disabled circle,#mermaid-svg-hUkqAxA9OLdu2Wup .disabled text{fill:lightgray;}#mermaid-svg-hUkqAxA9OLdu2Wup .disabled text{fill:#efefef;}#mermaid-svg-hUkqAxA9OLdu2Wup .section-8 rect,#mermaid-svg-hUkqAxA9OLdu2Wup .section-8 path,#mermaid-svg-hUkqAxA9OLdu2Wup .section-8 circle,#mermaid-svg-hUkqAxA9OLdu2Wup .section-8 polygon,#mermaid-svg-hUkqAxA9OLdu2Wup .section-8 path{fill:hsl(150, 100%, 76.2745098039%);}#mermaid-svg-hUkqAxA9OLdu2Wup .section-8 text{fill:black;}#mermaid-svg-hUkqAxA9OLdu2Wup .node-icon-8{font-size:40px;color:black;}#mermaid-svg-hUkqAxA9OLdu2Wup .section-edge-8{stroke:hsl(150, 100%, 76.2745098039%);}#mermaid-svg-hUkqAxA9OLdu2Wup .edge-depth-8{stroke-width:-10;}#mermaid-svg-hUkqAxA9OLdu2Wup .section-8 line{stroke:hsl(330, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-hUkqAxA9OLdu2Wup .disabled,#mermaid-svg-hUkqAxA9OLdu2Wup .disabled circle,#mermaid-svg-hUkqAxA9OLdu2Wup .disabled text{fill:lightgray;}#mermaid-svg-hUkqAxA9OLdu2Wup .disabled text{fill:#efefef;}#mermaid-svg-hUkqAxA9OLdu2Wup .section-9 rect,#mermaid-svg-hUkqAxA9OLdu2Wup .section-9 path,#mermaid-svg-hUkqAxA9OLdu2Wup .section-9 circle,#mermaid-svg-hUkqAxA9OLdu2Wup .section-9 polygon,#mermaid-svg-hUkqAxA9OLdu2Wup .section-9 path{fill:hsl(180, 100%, 76.2745098039%);}#mermaid-svg-hUkqAxA9OLdu2Wup .section-9 text{fill:black;}#mermaid-svg-hUkqAxA9OLdu2Wup .node-icon-9{font-size:40px;color:black;}#mermaid-svg-hUkqAxA9OLdu2Wup .section-edge-9{stroke:hsl(180, 100%, 76.2745098039%);}#mermaid-svg-hUkqAxA9OLdu2Wup .edge-depth-9{stroke-width:-13;}#mermaid-svg-hUkqAxA9OLdu2Wup .section-9 line{stroke:hsl(0, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-hUkqAxA9OLdu2Wup .disabled,#mermaid-svg-hUkqAxA9OLdu2Wup .disabled circle,#mermaid-svg-hUkqAxA9OLdu2Wup .disabled text{fill:lightgray;}#mermaid-svg-hUkqAxA9OLdu2Wup .disabled text{fill:#efefef;}#mermaid-svg-hUkqAxA9OLdu2Wup .section-10 rect,#mermaid-svg-hUkqAxA9OLdu2Wup .section-10 path,#mermaid-svg-hUkqAxA9OLdu2Wup .section-10 circle,#mermaid-svg-hUkqAxA9OLdu2Wup .section-10 polygon,#mermaid-svg-hUkqAxA9OLdu2Wup .section-10 path{fill:hsl(210, 100%, 76.2745098039%);}#mermaid-svg-hUkqAxA9OLdu2Wup .section-10 text{fill:black;}#mermaid-svg-hUkqAxA9OLdu2Wup .node-icon-10{font-size:40px;color:black;}#mermaid-svg-hUkqAxA9OLdu2Wup .section-edge-10{stroke:hsl(210, 100%, 76.2745098039%);}#mermaid-svg-hUkqAxA9OLdu2Wup .edge-depth-10{stroke-width:-16;}#mermaid-svg-hUkqAxA9OLdu2Wup .section-10 line{stroke:hsl(30, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-hUkqAxA9OLdu2Wup .disabled,#mermaid-svg-hUkqAxA9OLdu2Wup .disabled circle,#mermaid-svg-hUkqAxA9OLdu2Wup .disabled text{fill:lightgray;}#mermaid-svg-hUkqAxA9OLdu2Wup .disabled text{fill:#efefef;}#mermaid-svg-hUkqAxA9OLdu2Wup .section-root rect,#mermaid-svg-hUkqAxA9OLdu2Wup .section-root path,#mermaid-svg-hUkqAxA9OLdu2Wup .section-root circle,#mermaid-svg-hUkqAxA9OLdu2Wup .section-root polygon{fill:hsl(240, 100%, 46.2745098039%);}#mermaid-svg-hUkqAxA9OLdu2Wup .section-root text{fill:#ffffff;}#mermaid-svg-hUkqAxA9OLdu2Wup .section-root span{color:#ffffff;}#mermaid-svg-hUkqAxA9OLdu2Wup .section-2 span{color:#ffffff;}#mermaid-svg-hUkqAxA9OLdu2Wup .icon-container{height:100%;display:flex;justify-content:center;align-items:center;}#mermaid-svg-hUkqAxA9OLdu2Wup .edge{fill:none;}#mermaid-svg-hUkqAxA9OLdu2Wup .mindmap-node-label{dy:1em;alignment-baseline:middle;text-anchor:middle;dominant-baseline:middle;text-align:center;}#mermaid-svg-hUkqAxA9OLdu2Wup :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} 训练一个 GPT
数据
字符级 Tokenizer
整体放进显存
标签右移一位
模型
QKV 合并投影
SDPA 三种分支
Weight Tying
GPT-2 初始化
训练
Warmup 加 Cosine
AdamW 分组衰减
bf16 TF32 compile
按 val loss 早停
生成
Greedy 会重复
Temperature
Top-k
KV Cache
Prefill 与 Decode
结果逐字一致
小模型 launch-bound
大 batch 长序列加速 9 倍

一句话总结：GPT 就是"一个 Causal Decoder + 预测下一个 Token"。训练时靠 Causal Mask 并行算出所有位置的 loss，靠早停对抗小数据上的过拟合；推理时靠 KV Cache 避免重复计算------但它只有在计算量真正成为瓶颈时才显出价值。

本章思考题

step 0 的 loss（4.30）为什么接近但略高于 ln⁡65\ln 65ln65？如果 Embedding 按 N(0,1)\mathcal{N}(0, 1)N(0,1) 初始化，初始 loss 会更高还是更低？
为什么 train loss 和 val loss 的 gap 会单调增大？如果 val loss 一直在降、gap 却在变大，模型是在过拟合吗？
KV Cache 的显存占用如何计算？对于 32 层、d_model=4096、上下文 4096 的模型，单条请求的 KV Cache 有多大（FP16）？GQA（8 个 KV 头，32 个 Q 头）能把它降低多少？
为什么训练时不需要 KV Cache？
21.9 节中，batch 从 1 增加到 32，KV Cache 生成 256 个 Token 的耗时几乎不变。这对部署推理服务意味着什么？

下一章预告

GPT 能写出"像莎士比亚"的句子，但它不知道任何训练数据以外的事实，而且会一本正经地编造。下一章把前面的知识串成一个完整的应用：实现一个完整的 RAG 系统。

用本书前 19 章做知识库，构建本地问答系统
结构感知分块、混合检索（向量 + BM25）、Reranker 精排，全部跑在 GPU 上
用 transformers 直接加载 Qwen2.5-7B-Instruct（bf16）生成带引用的回答
用 Recall@K / MRR 量化每个组件带来的提升，用 Reranker 分数在检索层拒答

第21章 训练一个简单的 GPT