跳到主要内容

数据截至 (上游 commit 3adf61e154c3)

01 · GPT 模型：一个文件讲完 Transformer

这一章讲什么： model.py 的 331 行如何装下一个完整的 GPT——从注意力、残差块这些「骨架」，到初始化、权重共享、加载外部权重这些「不写在教科书公式里的决定」。读完你应该能不看注释复写出这个文件的大纲。

1. 它要解决的小问题​

「一个 GPT 模型」在论文里是一句话：堆 N 层 Transformer block，最后接 softmax 预测下一个 token。但要变成能跑、能训得动、还能和 OpenAI 官方 checkpoint 对齐的代码，必须回答一串论文不答的问题：

  • Q/K/V 怎么算才快？（答案：一次矩阵乘出全部，再切开）
  • 几十层堆起来梯度不炸的诀窍是什么？（答案：pre-LN + 特殊初始化）
  • 参数怎么初始化？（答案：正态 0.02，残差投影再缩 1/√(2L)）
  • 别人的 checkpoint 怎么搬进来？（答案：逐 key 拷贝，四个矩阵要转置）

model.py 就是这四个答案的合集。

2. 顶层结构：零件清单​

GPT.__init__（model.py:120-148）把零件全摆出来，结构只有三层嵌套：

GPT
├── transformer.wte   token 嵌入 (vocab_size × n_embd)
├── transformer.wpe   位置嵌入 (block_size × n_embd)
├── transformer.drop  dropout
├── transformer.h     n_layer 个 Block
│     └── 每个 Block: ln_1 → attn → ln_2 → mlp
├── transformer.ln_f  最终 LayerNorm
└── lm_head           输出投影 (n_embd × vocab_size, 无 bias)

前向（GPT.forward，model.py:170-193）就是把这张清单顺序走一遍：token 嵌入加位置嵌入 → dropout → 逐个 Block → 最终 LayerNorm → lm_head 出 logits。

零件类在哪
层归一化LayerNorm（支持可选 bias，PyTorch 原生不支持 bias=False）model.py:18-27
因果自注意力CausalSelfAttentionmodel.py:29-76
前馈网络MLPmodel.py:78-92
一个残差块Blockmodel.py:94-106
超参包GPTConfig dataclassmodel.py:108-116
总装 + 前向GPTmodel.py:118-193

3. 机制一：Block 与 pre-LN 残差​

小问题： Transformer 块有两个子层（注意力、MLP），LayerNorm 和残差连接怎么摆，直接决定深网络能不能训动。

直觉： 现代做法（GPT-2 沿用）是 pre-LN——先归一化再进子层，子层输出直接加回主干。残差通路全程是「干净的高速公路」，梯度可以一路传回第一层，不被 LayerNorm 反复缩放。

真实实现短到可以全引（Block.forward，model.py:103-106）：

def forward(self, x):
    x = x + self.attn(self.ln_1(x))
    x = x + self.mlp(self.ln_2(x))
    return x

重点看：两个 x + 的左操作数都是未经任何变换的 x——这就是「高速公路」。LayerNorm 只出现在支路上。

坑/细节： 这个 LayerNorm 是自写的（model.py:18-27），唯一原因是 PyTorch 原生的 nn.LayerNorm 不让关 bias；GPT-2 有 bias，但 nanoGPT 训练时默认 bias=False（train.py:56），快一点也略好一点。

4. 机制二：CausalSelfAttention——因果自注意力​

小问题： 每个位置要「看左边所有位置、加权混合它们的信息」，且绝不能偷看右边（否则预测下一个 token 就等于作弊看答案）。

直觉： 三投影出 Q/K/V，用 QKᵀ/√d 算相似度，softmax 归一成权重，再加权求和 V。「因果」靠一个下三角掩码把右上角的相似度全部置 -∞。

图示（单头的数据形状流）：

x (B, T, C)
   │  c_attn 一次矩阵乘出 3C
   ▼
q, k, v 各 (B, T, C) ──► reshape 成 (B, nh, T, hs)
   │
   ▼  q @ kᵀ / √hs   →   (B, nh, T, T) 注意力分数
掩码: 上三角 = -∞
   │
   ▼  softmax → @ v   →   (B, nh, T, hs)
拼回头 → c_proj → (B, T, C)

真实实现的三个关键点：

  1. QKV 一次算完再切。 c_attn 是一个 n_embd → 3*n_embd 的 Linear（model.py:35），前向里 .split(self.n_embd, dim=2) 切成 q/k/v（model.py:56）——三次矩阵乘合成一次，GPU 更友好。
  2. 有 Flash Attention 就用。 __init__ 探测 torch.nn.functional.scaled_dot_product_attention 是否存在（model.py:45）；存在则一行调用、靠 is_causal=True 保证因果（model.py:64）。
  3. 没有就走手写路径。 显式 QKᵀ、用注册的下三角 buffer 做掩码、softmax、乘 V（model.py:67-71）。注意这个掩码 buffer 的名字叫 bias，是为了和 HF GPT-2 的 state_dict 对齐（model.py:49-50）。

坑： 掩码 buffer 只有在 self.flash 为 False 时才注册（model.py:46-50）。这导致一个有 Flash 的环境存出的 state_dict 里没有 attn.bias 这个 key——from_pretrained 和 crop_block_size 都要专门处理它（后者在 model.py:203-204 用 hasattr 防御）。

5. 机制三：权重初始化——让深网络从第一步就能训​

小问题： 12 层（甚至 48 层）残差块串起来，如果每层残差分支的输出方差都差不多大，主干信号的方差会随层数不断累加，初始化时数值就漂了。

直觉： 让所有「汇入残差主干」的最后一层投影（即每个 attn 和 mlp 里的 c_proj）的初始权重按层数缩小：标准差除以 √(2·n_layer)（每个 block 有 2 个残差分支，方差叠加按平方根缩放）。

真实实现（model.py:140-145）：

self.apply(self._init_weights)  # 全部 Linear/Embedding: normal(0, 0.02)
for pn, p in self.named_parameters():
    if pn.endswith('c_proj.weight'):
        torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer))

顺序很关键：先 _init_weights（model.py:162-168）把所有 Linear/Embedding 统一成 normal(0, 0.02)，再按参数名后缀把 c_proj 单独覆盖成更小的 std。注释明说这来自 GPT-2 论文（model.py:142）。

6. 机制四：权重共享（weight tying）​

小问题： token 嵌入矩阵 wte（vocab×n_embd）和输出投影 lm_head（n_embd×vocab）形状互为转置——一个把 id 变成向量，一个把向量变成 id 上的打分。要不要用两份参数？

nanoGPT 的选择：一份。 直接赋值让两者共享同一块内存（model.py:138）：

self.transformer.wte.weight = self.lm_head.weight

好处是省掉 vocab×n_embd 个参数（GPT-2 规模下约 3850 万，占 124M 的将近三分之一），实践上效果不差甚至更好（注释链到 paperswithcode 的 weight tying 词条，model.py:138）。

细节： get_num_params 里有一段专门解释参数计数口径——位置嵌入不计入，token 嵌入因为被 lm_head 复用所以计入（model.py:150-160）。报参数量时这是个常见的对不上账的地方。

7. 机制五：from_pretrained——把 OpenAI 的权重原样搬进来​

小问题： 想微调 GPT-2，就得把 OpenAI 当年发布的 checkpoint 装进这个模型。但两边不只是文件格式不同，连权重的摆放方式都不同。

两个真实障碍，代码各自给出一行答案：

  1. Conv1D vs Linear。 OpenAI 的 checkpoint 里四个矩阵（c_attn、c_proj、mlp.c_fc、mlp.c_proj）是以 Conv1D 形式存的，形状和 Linear 权重互为转置。搬的时候这四个 key 要 .t() 转置后拷贝，其余 key 直接拷（model.py:245-259，transposed 列表 + sd[k].copy_(sd_hf[k].t())）。
  2. 掩码 buffer 不算参数。 HF 的 state_dict 里带着 attn.bias / attn.masked_bias 这种 buffer，不是参数，拷贝前要过滤掉（model.py:235 和 model.py:243-244），否则 key 数对不上。

整个搬运用「两边 key 排序后等长、逐个断言 shape 匹配」来保证正确性（model.py:248、model.py:252、model.py:257）——任何一个 key 对不上都会立刻 assert 炸出来，而不是静默错训。

坑： crop_block_size（model.py:195-204）能把 1024 的位置嵌入手术裁短，方便在小上下文上微调；但它只裁 wpe 和手写路径的掩码 buffer，Flash 路径不受影响（掩码根本不存在）。这也是「attn.bias 时有时无」这一坑的另一面。

8. 前向里的两个不显然决定​

两个容易扫过去、其实很讲究的细节：

  • 训练/推理共用一个 forward，靠 targets is None 分叉。 训练时算全序列的交叉熵，ignore_index=-1 允许调用方掩盖不想要的位置（model.py:184-187）；推理时只对最后一个位置过 lm_head，省掉 T-1 倍的输出投影计算（model.py:189-191）。注释原话：「inference-time mini-optimization」。
  • 序列超长直接 assert。 t <= block_size 硬检查（model.py:173），不静默截断——训练侧出错应该立刻响，而不是悄悄训出坏模型。

9. 本章代码地图​

主题文件符号
可选 bias 的 LayerNormmodel.py:18-27LayerNorm
因果自注意力model.py:29-76CausalSelfAttention
Flash/手写双路径model.py:62-71CausalSelfAttention.forward
前馈网络model.py:78-92MLP
pre-LN 残差块model.py:94-106Block
超参 dataclassmodel.py:108-116GPTConfig
模型总装model.py:120-148GPT.__init__
初始化model.py:162-168、model.py:143-145_init_weights、c_proj 缩放
权重共享model.py:138wte.weight = lm_head.weight
前向model.py:170-193GPT.forward
上下文裁短model.py:195-204crop_block_size
搬 GPT-2 权重model.py:206-261GPT.from_pretrained
优化器分组model.py:263-287configure_optimizers（下一章细讲）
生成model.py:305-330generate（第 4 章细讲）