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；attention 不知道顺序，位置必须显式给

GPT-2 small 的 decoder-only 栈：embedding + 位置表 → 12 个 block（causal attention + FFN）→ lm_head

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

encoder 双向读源句、decoder causal 生成、cross-attention 读 encoder 输出；GPT-2 走 decoder-only。

encoder-decoder 与 decoder-only 的两条路线对比

一个 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 的旅程：训练侧与推理侧

训练 / prefill / decode 是三种形态——KV cache = 算一次存下来

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

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 逐行

结构本身不到 90 行；对拍 HF GPT-2 相对差 9e-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 与训一个会续写的模型

0.8M 参数 2000 步 7 分钟，loss 4.17 → 1.66

只改层数在 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