Transformer 原理与实现：从 GPT-2 总览到手搓 GPT-2

系列精华 · 四篇正文，按 ↓ 看推导、代码与数字

arganzheng · 2026-04-09

这个系列的一句话主张

Transformer 的结构可以从五种运算读懂；沿着一个 token 的旅程，再把同一套结构写成代码并训练起来。

篇 内容
01–04 静态结构 → token 旅程 → nanoGPT model.py → train.py 与实际训练

四篇怎么连起来

flowchart LR
    P1["01 静态结构"] --> P2["02 token 的旅程"] --> P3["03 model.py"] --> P4["04 train.py"]

01 · GPT-2 的 decoder-only 总览

先看本系列要拆解的 GPT-2 small 总览；接着回到 2017 年的论文，解释 encoder-decoder、三条路线与 GPT-2 的来历。

结论：GPT-2 走 decoder-only 路线。embedding 查表 → attention（唯一让 token 互相看的地方） → FFN（逐 token 的非线性，知识在这，占一层 2/3）→ 残差 + LayerNorm → lm_head；attention 不知道顺序，位置必须显式给。

%%{init: {"flowchart": {"wrappingWidth": 330}}}%%
flowchart LR
    IN["token 编号<br/>[464, 3797, 3332, 319, 262]"] --> EMB["embedding 查表<br/>50257 × 768 → 5 个 768 维向量<br/>+ 位置表"]
    EMB --> B
    subgraph B["× 12 个 block"]
        direction TB
        A["LayerNorm → causal self-attention → 加回"]
        F["LayerNorm → FFN 768→3072→768 → 加回"]
        A --> F
    end
    B --> HEAD["LayerNorm → lm_head 768 × 50257<br/>softmax → p(mat) = 0.31"]

01 · 回到原始 Transformer：encoder-decoder 的两半

《Attention Is All You Need》里的 Transformer 是为机器翻译设计的 encoder-decoder：encoder 双向读取源句，decoder 用 causal self-attention 生成目标句，再以 cross-attention 读取 encoder 输出。GPT-2 属于后来常见的 decoder-only 路线（原始论文）。

flowchart LR
    X["源句"] --> E["Encoder<br/>双向 self-attention + FFN"]
    E --> M["Encoder 输出"]
    Y["已生成目标前缀"] --> D["Decoder<br/>causal self-attention<br/>cross-attention + FFN"]
    M -->|K、V| D
    D --> P["下一个目标 token"]

一个 attention 头的六步手算（d = 4、T = 3）

① x 乘三个矩阵得到 Q、K、V；② S = QKᵀ；③ 除 √d；④ 加 causal mask；⑤ softmax；⑥ 乘 V

  • \(\text{Attention}(Q,K,V) = \text{softmax}(QK^\top/\sqrt d + M)\,V\)；t2 的权重 (0.27, 0.27, 0.45)，与 PyTorch 对拍差 \(6 \times 10^{-8}\)
  • 随机 \(q \cdot k\) 的标准差是 \(\sqrt d\)：不除，d = 128 时 softmax 一个 token 独占

GPT-2 真实的头在看什么

GPT-2 small 处理 The cat sat on the mat because it was tired：第 4 层第 11 头看上一个词，第 3 头把 it 指回 cat（0.84）

  • 多头不增加算量：h 个小乘法 = 一个大乘法，只多一个 \(W_O\)
  • 换序实验：把输入打乱，attention 输出跟着换位——所以位置必须另给（GPT-2 查表 / Llama RoPE）

02 · 一个 token 的旅程：训练侧与推理侧

结论：causal mask + teacher forcing 让一句话 = T 个训练样本；推理时旧 token 的 K、V 不变，算一次存下来就是 KV cache——训练 / prefill / decode 是三种形态。

训练的一步：五个位置一次前向同时得到五个分布，目标是输入右移一位，五个交叉熵取平均

prefill 与 decode：为什么前面的 token 不用重算

prefill 把五个 token 一次算完只取最后位置，K、V 存进 cache；decode 每步只算新 token 的 Q、K、V 并追加

形态 输入 有反向 存什么 瓶颈
训练 [B, T] 全位置有用 是 激活值 算力
prefill [B, T] 只取最后 否 K、V 算力（TTFT）
decode [B, 1] 否 追加 K、V 访存（TPOT）

数字

  • 有 / 无 KV cache 输出逐 token 一致；prompt 256 生成 256 快 7.9 倍
  • 代价：Llama-3-8B 128 KiB / token = 32 层 × 2 × 8 个 KV 头 × 128 × 2 B
  • KV cache 不存 Q——旧 token 的 Q 用过即弃
  • 训练长上下文比推理多一份与 B × T 成正比的激活值

03 · 手搓 GPT（上）：nanoGPT model.py 逐行

结论：330 行、6 个类，结构本身（LayerNorm、CausalSelfAttention、MLP、Block）不到 90 行；与 HF GPT-2 对拍相对差 \(9 \times 10^{-5}\)。

class Block(nn.Module):
    def forward(self, x):
        x = x + self.attn(self.ln_1(x))   # pre-norm：先 norm 再子层，残差直连
        x = x + self.mlp(self.ln_2(x))
        return x

# CausalSelfAttention.forward 的骨架
q, k, v = self.c_attn(x).split(self.n_embd, dim=2)          # 一次 GEMM 算 QKV
k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)  # 头换到第 1 维
y = F.scaled_dot_product_attention(q, k, v, is_causal=True)   # 融合 kernel
y = y.transpose(1, 2).contiguous().view(B, T, C)              # 拼回去

每个类落在结构的哪个方框

类 行数 实现的是
LayerNorm 10 可选 bias（GPT-2 有、Llama 无）
CausalSelfAttention 40 c_attn 一次算 QKV · 拆头 · SDPA · c_proj
MLP 12 768 → 3072 → GELU → 768
Block 10 上面两行 pre-norm
GPT 200+ 拼结构、共享权重、初始化、generate、from_pretrained、configure_optimizers、estimate_mfu
  • Llama 相对 GPT-2 只改五处：RMSNorm、RoPE、SwiGLU、GQA、去 bias（lm_head 不共享）
  • 三个独立 Linear 与一个 c_attn 数学等价，合并只为一次 GEMM

04 · 手搓 GPT（下）：train.py 与训一个会续写的模型

结论：1.1 MB 莎士比亚、0.8M 参数、2000 步 7 分钟，loss 从 ln 65 = 4.17 → 1.66；层是串行的——参数与耗时线性、loss 收益递减。

只改层数在 shakespeare_char 上 2000 步的 loss：2 / 4 / 8 层 val 1.82 / 1.66 / 1.59，每步耗时约 0.5× / 1× / 2×

训练循环每一步在做什么

%%{init: {"flowchart": {"wrappingWidth": 170}}}%%
flowchart LR
    LR["设 lr<br/>warmup + cosine"] --> EV["到点评估<br/>存 checkpoint 五样"] --> ACC["k 个 micro-batch<br/>loss ÷ k 累积"] --> CL["unscale<br/>+ 梯度裁剪"] --> ST["optimizer.step()<br/>zero_grad()"] --> LR
  • get_batch：memmap + 随机窗口 + stack，y 是 x 右移一位，没有 epoch
  • 不除累积步数 ≠ 只是尺度不同：等价于学习率放大 k 倍；DDP 只在最后一个 micro-batch 同步
  • checkpoint 五样：model、optimizer、model_args、iter_num、config——续训不带优化器状态，前几百步抖
← Arganzheng's Blog