第 1 周 · Day 3：从零实现 mini-GPT，与 LSTM 做对照实验

时长 8 小时：§1–§2 搭模型（2h）· §3–§4 初始化与训练配方（1h）· 训练与对照实验（3h，边跑边读 §5–§7）· §A–§C（1h）· 读论文（1h，可挪到晚上）

对应正课 Lecture 1 的收尾：把整套架构落到代码上。学习计划里”本计划最重要的一个项目”就是今天。

配套代码 w1day3_checks.py（纯 numpy）：mini-GPT 与 LSTM 的精确参数量、用 numpy 前向一个刚初始化的 mini-GPT 验证初始 loss、训练预算与 6ND、学习率曲线、显存账、生成时的上下文裁剪

§B 参考实现（code/ 目录）：w1d3_minigpt.py（训练、采样、与 LSTM 对照、按位置的 loss）。先自己写，卡住超过 30 分钟再看。

今天的目标：一个自己写的、能训到验证集 BPC ≈ 2.1–2.3 的字符级 GPT；一份控制了变量的 GPT vs LSTM 对照结果；以及一个你能解释每一行的代码库，Day 4–5 会在它上面加 KV Cache、GQA 和 RoPE。

§0 今天的主线

0.1 为什么这是本周最重要的一天

前两天的每个结论，都是”在纸上成立”。今天要回答的是：把它们拼起来，真的能训出一个会写莎士比亚的模型吗？ 以及更难的一个问题：它比 LSTM 好，好在哪里，好多少？

Week 0 Day 5 §D.4 埋下了这个实验：同样的数据、同样的训练循环、同样的解码，只把中间那一层从 nn.LSTM 换成 Transformer decoder。这是一次干净的对照实验，也是整门课里最有说服力的一次演示。

0.2 对照实验的思想

科学实验的核心是只改变一个变量。如果 GPT 比 LSTM 好，你想知道”好”来自哪里：是架构本身？还是参数更多？还是学习率更合适？今天的代码刻意做了三件事来控制变量：

1. 数据管线、词表、切分完全复用 Week 0 Day 5（get_batch 一字不改）。
2. 训练循环共用一个函数：同样的优化器分组、调度、裁剪、评估窗口。
3. 额外训练一个参数量对齐的 LSTM，把”参数多”这个混杂因素分离出来（§5.2）。

0.3 今天要用到的旧结论

§1 从 CharLM 到 GPT

1.1 只换中间那层

                  Week 0 CharLM                     今天的 GPT
token ids (B,T)
   ↓        Embedding(V, d)                   Embedding(V, d)  +  位置嵌入(block, d)   ← 新增：Day 1 §1.5
   ↓        nn.LSTM(d, d, 2 层)               [ Pre-LN 块 ] × L                        ← 换掉：Day 1–2
   ↓        LayerNorm(d)                      LayerNorm(d)（末端 LN）
   ↓        Linear(d, V)，与嵌入绑定            Linear(d, V)，与嵌入绑定
logits (B,T,V)

只多了一样东西：位置嵌入。 原因在 Day 1 §1.5：LSTM 按顺序读，天然知道位置；自注意力置换等变，必须外挂。其余部分，包括损失、采样、评估、checkpoint，全部和 Week 0 一样。

1.2 配置与参数量

默认配置沿用 nanoGPT 的 Shakespeare 字符级示例，方便对照公开结果：

w1day3_checks.py 第 1 组按 Day 2 §5.1 的公式精确数出 10,770,816。其中非嵌入部分 10.62M，词嵌入只有 $65\times384=24{,}960$：字符级词表很小，参数几乎全在块里，和 GPT-2 small 词表占 31% 的情况正好相反（Day 2 §6.4）。

Week 0 的 CharLM（2 层 LSTM，$d=384$）只有 2.39M。默认配置下 GPT 的参数是它的约 4.5 倍。这一点在 §5.2 会变得重要。

§2 代码逐段

完整代码在 code/w1d3_minigpt.py。这里只讲每一段为什么这样写。

2.1 因果自注意力

class CausalSelfAttention(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.h, self.dk = cfg.n_head, cfg.d_model // cfg.n_head
        self.qkv  = nn.Linear(cfg.d_model, 3 * cfg.d_model, bias=cfg.bias)   # ①
        self.proj = nn.Linear(cfg.d_model, cfg.d_model, bias=cfg.bias)       # ②
        self.p = cfg.dropout
        self.resid_drop = nn.Dropout(cfg.dropout)

    def forward(self, x):
        B, T, d = x.shape
        q, k, v = self.qkv(x).split(d, dim=-1)
        q, k, v = (t.view(B, T, self.h, self.dk).transpose(1, 2) for t in (q, k, v))
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True,
                                           dropout_p=self.p if self.training else 0.0)   # ③
        y = y.transpose(1, 2).contiguous().view(B, T, d)                    # ④
        return self.resid_drop(self.proj(y))

① 三个投影合成一个 Linear(d, 3d)：一次大矩阵乘比三次小的快，数学上完全等价（把 $W_Q,W_K,W_V$ 横向拼起来）。nn.MultiheadAttention 的 in_proj_weight 也是这么存的（Day 1 §A.2）。

② 输出投影叫 proj，初始化时要按 $1/\sqrt{2L}$ 缩放（§3.3），代码里按名字找它。

③ 训练时用 SDPA 的 is_causal=True：这里 $L=S=T$，左上角对齐的下三角就是正确的因果掩码。Day 1 §A.1 的边界 3（$L\ne S$ 时出错）要到 Day 4 的增量解码才会碰到。dropout_p 要自己看 self.training（边界 2）。

④ 合并多头：先 transpose 再 view，顺序不能反（Day 1 §A.5）。

代码里还保留了一个 attn_impl="manual" 分支，用 Day 1 的手写版本，自检会验证两者逐元素一致。调试时切到 manual 能拿到注意力矩阵。

2.2 FFN 与块

class Block(nn.Module):
    def forward(self, x):
        x = x + self.attn(self.ln1(x))      # 通信
        x = x + self.mlp(self.ln2(x))       # 计算
        return x

就是 Day 2 §2.1 那张图。MLP 默认用 GELU 的 tanh 近似（和 GPT-2 一致，Day 2 §A.2），可选 ReLU / SwiGLU。--post-ln 切换成 Post-LN 做消融。

2.3 GPT 主体

class GPT(nn.Module):
    def __init__(self, cfg):
        self.wte = nn.Embedding(cfg.vocab_size, cfg.d_model)
        self.wpe = nn.Embedding(cfg.block_size, cfg.d_model)      # 可学习位置嵌入（Day 2 §3.7）
        self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layer)])
        self.ln_f = nn.LayerNorm(cfg.d_model)                     # Pre-LN 必须有末端 LN（Week 0 Day 2 §4.5）
        self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
        self.apply(self._init)                                    # 先初始化……
        ...                                                       # 残差分支 × 1/√(2L)
        self.lm_head.weight = self.wte.weight                     # ……再绑定

    def forward(self, idx, targets=None):
        pos = torch.arange(idx.size(1), device=idx.device)       # device 别忘（Day 2 §A.3）
        x = self.drop(self.wte(idx) + self.wpe(pos))
        for b in self.blocks:
            x = b(x)
        logits = self.lm_head(self.ln_f(x))
        ...

为什么 ModuleList 而不是 list？ Week 0 Day 1 §A.1：普通 list 里的模块不会被注册，参数不会进优化器，不报错，只是那些层永远不训练。

lm_head 没有 bias：绑定的是权重，bias 没有对应物；GPT-2 也没有。

2.4 生成：裁剪上下文

@torch.no_grad()
def generate(self, idx, max_new, temperature=1.0, top_k=None, top_p=None):
    for _ in range(max_new):
        ctx = idx[:, -self.cfg.block_size:]          # 只保留最后 block_size 个
        logits, _ = self(ctx)
        nxt = sample_next(logits[:, -1, :], temperature, top_k, top_p)   # Week 0 Day 5 §5
        idx = torch.cat([idx, nxt], dim=1)
    return idx

裁剪是必须的：位置嵌入表只有 block_size 行，第 257 个位置没有向量可查（w1day3_checks.py 第 6 组）。LSTM 没有这个限制，因为它不需要位置嵌入。这是绝对位置编码的代价，Day 5 的 RoPE 会部分缓解它。

这个生成循环每一步都把整个上下文重新前向一遍。 生成第 257 个字符时，前 255 个字符的 K、V 其实上一步已经算过了，这里又算了一遍。这正是 Day 4 KV Cache 要消除的浪费；Week 0 Day 5 §5.1 的 LSTM 反而天然是增量的。

§3 初始化

3.1 为什么是 $\mathcal N(0,0.02)$

GPT-2 把所有线性层和嵌入都初始化成标准差 0.02。它不是从 Week 0 Day 2 §5.1 的二阶矩传播推出来的（那个推导给出 $1/\sqrt{\text{fan\_in}}$，$d=384$ 时约 0.051），而是一个经验值。能工作的原因是 Pre-LN：每个子层的输入都先被 LayerNorm 归一化，残差流的尺度由 LN 控制，不依赖每层的初始化恰好保持方差。

3.2 验证：初始 loss 应该是多少

Week 0 Day 5 §4.3 说过：初始 loss 应 $\approx\ln V$。现在可以算得更精确。

末端 LN 输出每维方差约 1，乘上标准差 0.02 的绑定权重，logits 的标准差约 $0.02\sqrt d=0.392$（$d=384$）。logits 有小的随机波动时，$\mathbb E[\log\sum_je^{z_j}]\approx\log V+\sigma^2/2$，于是

$\mathcal L_0\approx\ln65+\frac{0.392^2}{2}=4.174+0.077=4.251$

w1day3_checks.py 第 2 组用 numpy 完整前向了一个刚初始化的 6 层 mini-GPT（5 个种子）：logits 标准差实测 0.416，初始 loss 4.279 ± 0.023。比 $\ln V$ 高一点点，完全正常。如果你的初始 loss 明显高于 4.3，先查初始化；如果是 5、6 甚至更高，八成是某处用了默认初始化或忘了绑定后的缩放（Week 0 Day 5 §B 第 3 题）。

3.3 残差分支 $\times1/\sqrt{2L}$

Day 2 §2.3 实测过：不缩放时残差流的 RMS 按 $\sqrt L$ 增长。代码里：

for name, p in self.named_parameters():
    if name.endswith("proj.weight"):                     # attn.proj 和 mlp.proj
        nn.init.normal_(p, std=0.02 / math.sqrt(2 * cfg.n_layer))

6 层时影响不大，但层数一多就是稳定性问题。养成习惯。

3.4 为什么这里不像原始论文那样把嵌入乘 $\sqrt d$

Day 2 §3.8 说原始 Transformer 把词嵌入乘了 $\sqrt d$，因为它的位置编码是幅度 $O(1)$ 的正弦，词嵌入的初始化又很小。GPT-2 的位置嵌入是可学习的、和词嵌入同样用 0.02 初始化，两者尺度天然一致，所以不需要放大。用正弦位置编码（--pos sin）时，这个尺度问题会回来：0.02 量级的词嵌入加上 $O(1)$ 的位置编码，一开始位置信号压倒词义。这是 §5.4 的一个消融点。

§4 训练配方

4.1 每个超参的理由

$\beta_2$ 为什么是 0.99 而不是 LLM 常用的 0.95？ Week 0 Day 2 §6.3 说过，$\beta_2=0.95$ 让二阶矩”近视”，对梯度尖峰反应快，适合大 batch 的长时间训练。这里每步只有 16k token，梯度噪声大，二阶矩估计需要更长的平均窗口（$1/(1-\beta_2)$：0.99 约 100 步，0.95 约 20 步）才稳定。这是一个取舍，不是定律；可以作为消融。

学习率曲线（w1day3_checks.py 第 4 组）：

4.2 预算：要跑多久

w1day3_checks.py 第 3 组。默认配置每步 $64\times256=16{,}384$ 个 token，5000 步共 81.9M token，是训练集（1,003,854 个字符）的 81.6 遍。训练计算量 $6ND=6\times10.77\text{M}\times81.9\text{M}=5.3\times10^{15}$ FLOPs：

“实际吞吐”远低于硬件的峰值，小模型尤其如此（kernel 启动、数据搬运的开销占比大）。第一次跑一定先用 --small，几分钟内确认 loss 在降、生成在变好，再上默认配置。

过了 80 遍数据意味着什么？ 必然过拟合。这就是 dropout 0.2 和”按验证 loss 保存最好的 checkpoint”（早停）的原因。

4.3 显存

第 5 组（bf16 混合精度，用 Korthikanti et al. 2022 的激活公式）：

激活是参数状态的十几倍。 显存不够时减 batch 或 block_size。用 SDPA 的 FlashAttention 内核可以省掉注意力矩阵那一块（Day 6 §3）。

4.4 该看到什么

• 第 0 步：验证 loss ≈ 4.28（§3.2）。
• 前几百步：快速降到 2.5 左右，生成的字符频率对了（Week 0 Day 5 §4.3 的”早期”）。
• 1000–2000 步：拼出真词，出现人名 + 冒号 + 换行的剧本格式。
• 之后：验证 loss 趋于平台，训练 loss 继续下降，两者的差距拉大 = 过拟合。

公开参照：nanoGPT 在同样配置下报告的验证 loss 约 1.47（BPC ≈ 2.1），这是一个合理的目标。这些数字不是本仓库实测的，你跑出来的结果请记录下来，和 Week 0 的 LSTM 一起填进 §5 的表。

§5 对照实验

5.1 固定了什么

python w1d3_minigpt.py --compare-lstm           # 默认配置
python w1d3_minigpt.py --compare-lstm --small   # 小配置

它依次训练三个模型：GPT、宽度与 GPT 相同的 LSTM（$d=384$）、参数量对齐的 LSTM（$d=768$）。三者共用：同一份数据和切分、同一个 get_batch、同样的 batch 与 block、同样的步数（因此同样的 token 数）、同一个训练函数（同样的优化器分组、调度形状、裁剪、评估窗口）。不同的只有模型本身和学习率（每个模型用自己合适的学习率，否则比较的是”谁更适应这个学习率”）。

5.2 混杂因素：参数量

w1day3_checks.py 第 1 组给出 2 层 LSTM 在不同宽度下的参数量：

GPT（10.77M）和 $d=768$ 的 LSTM（9.50M）接近。如果 GPT 只赢过 LSTM-384 而输给或持平 LSTM-768，结论就是”参数多”而不是”架构好”。 这就是为什么要训第三个模型。

还有两个维度值得记录：同样的 token 数（上面固定了）和同样的墙钟时间。LSTM 沿时间串行，GPU 上每步更慢（Day 1 §5.3）；--compare-lstm 会打印用时。如果时间预算固定，GPT 能在同样时间里看更多数据。

5.3 更细的问题：模型在哪个位置更好？

平均 loss 只是一个数。--per-position 打印验证集上每个上下文位置的平均 loss：第 $t$ 个位置意味着模型已经看到了 $t+1$ 个字符。

你应该预期看到（请用实验检验，不要把它当结论）：

• $t=0$：只看到一个字符，模型此时就是一个 bigram 模型（Week 0 Day 5 §1.4 的小实验里 bigram 约 3.6 BPC；Shakespeare 上的具体数值请自己测），两个模型应该差不多。
• $t$ 增大：loss 下降，因为上下文提供了越来越多的信息。
• 两个模型的曲线在哪个区间分开？LSTM 的记忆要经过每一步的状态更新传递；注意力对上下文里的每个位置都是一步直达。

画出两条曲线，写下你观察到的差异和你的解释。这个图比平均 loss 更能说明”注意力到底带来了什么”。

5.4 消融清单

每个消融只改一处，其他全用默认。先写下你的预测，再跑：

§6 调试清单

先跑 --selftest：它自动检查下表里能自动检查的几条。

§7 读论文：Attention Is All You Need

学习计划要求”逐句读，读三遍”。有了前三天的准备，每一遍都有明确的任务。

第一遍（30 分钟）：只看结构。 读摘要、引言、结论，看 Figure 1 和 Figure 2。目标：用一句话说出这篇论文的贡献，并在图上指出 Day 1–2 讲过的每一个零件。

第二遍（1.5 小时）：逐节对照。

第三遍（1 小时）：Table 3 的消融。 这是全文最有信息量的表，每一行都是一次控制变量实验。重点看：(A) 行，单头比最好的设置差 0.9 BLEU，但头太多也会变差，对照 Day 1 §4.4 的低秩瓶颈；(B) 行，减小 $d_k$ 会伤害质量；(E) 行，可学习位置编码与正弦几乎一样。然后问自己：今天 §5.4 的消融，和他们的消融有哪些对应？

计划中另外两篇（Word2Vec、Bahdanau）在 Week 0 Day 3、Day 4 已经讲过核心内容，扫读即可：看它们的实验怎么设计，这是写自己实验报告时最好的模板。

§A 训练工程：API 用法与边界

A.1 torch.compile

model = torch.compile(model)          # PyTorch 2.x

• 第一次调用会编译，可能要几十秒，之后的步才快。计时要跳过前几步。
• 在 CUDA 上收益最明显；MPS 和 CPU 上的支持和收益因版本而异，不行就去掉。
• 改变了输入形状（比如生成时序列长度一直在变）会触发重新编译。生成时用未编译的模型，或者只编译训练。
• 编译后的模型 state_dict 的键会带 _orig_mod. 前缀，保存时存 model._orig_mod.state_dict()。

A.2 混合精度

torch.set_float32_matmul_precision("high")                 # Ampere 及以上启用 TF32，几乎白送的加速
with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
    logits, loss = model(x, y)
loss.backward()                                            # 在 autocast 之外

• bf16 不需要 GradScaler，fp16 需要（Week 0 Day 2 §A.5、Week 0 Day 1 §10.2）。T4 这类老卡不支持 bf16，只能 fp16 + scaler。
• backward() 放在 autocast 上下文之外。
• 损失计算在 autocast 里时，F.cross_entropy 会自动用 fp32 算，不用手动转。

A.3 生成时的上下文

• @torch.no_grad() 或 torch.inference_mode() 必加；model.eval() 必加（Week 0 Day 5 §A.5）。
• 超过 block_size 必须裁剪（§2.4）。
• torch.multinomial 在 MPS 上的支持和随机数行为与 CPU 不同，要可复现就在 CPU 上采样（代码里 sample_next 就是这么做的）。

A.4 checkpoint

保存时带上 config 和词表（Week 0 Day 5 §A.6）：

torch.save({"model": model.state_dict(), "config": asdict(cfg), "chars": vocab.chars}, path)

用 dataclass 存配置，加载时 GPTConfig(**ck["config"]) 一行重建模型。之后几天给配置加字段（n_kv_head、pos="rope"）时，给新字段设默认值，老 checkpoint 就还能加载。

§B 实操

1. 跑通 w1day3_checks.py，把第 1、2、3 组的数字和 §1.2、§3.2、§4.2 对上。

2. 自己写 mini-GPT（3h，不看参考代码）：从 Week 0 Day 5 的 CharLM 出发，只换中间那层。写完先跑三个检查：(a) 参数量等于公式；(b) 初始 loss ≈ 4.28；(c) 因果性的梯度检查。三个都过了再训练。

3. 训练并记录（边跑边读 §5–§7）：先 --small 确认流程，再跑默认配置。每 250 步记录训练 / 验证 loss，每 500 步读一段生成文本，对照 §4.4。

4. 对照实验：--compare-lstm，填下面这张表，再用 --per-position 画出两条位置曲线，写一段 200 字以内的结论（必须回应 §5.2 的参数混杂）。

   模型	参数	最好验证 loss	BPC	用时
   GPT	10.77M
   LSTM-384	2.39M
   LSTM-768	9.50M

5. 消融（挑 3 个，§5.4）：先写预测，再跑，再解释。

6. 读论文：按 §7 读完三遍，回答表里每一行的问题。

参考实现：code/w1d3_minigpt.py。跑完请把 --selftest 输出、训练日志和上面那张表发回来，文档里的预期会据此校准。

§C 自测

动机与直觉层

1. 从 CharLM 到 GPT，为什么只需要多加一个位置嵌入？
2. 一个”干净”的 GPT vs LSTM 对照实验需要固定哪些东西？为什么还要训一个加宽的 LSTM？
3. 为什么 GPT-2 用 $\mathcal N(0,0.02)$ 这种与宽度无关的初始化也能训起来？
4. 为什么这里 $\beta_2$ 用 0.99？它和 LLM 常用的 0.95 各自在权衡什么？
5. 为什么按位置的 loss 曲线比平均 loss 更能说明注意力的作用？

推导层

6. 推出默认配置的参数量 10,770,816。
7. 推出初始 loss $\approx\ln V+\sigma^2/2$ 中的 $\sigma$，并算出数值。
8. 算出默认配置的训练 token 数、过数据的遍数和 $6ND$。
9. 用 Korthikanti 的公式估算激活显存，说明哪一块能被 FlashAttention 省掉。

工程层

10. 为什么要先 apply(_init) 再绑定权重？反过来会怎样？
11. 生成时为什么要裁剪上下文？LSTM 为什么不用？
12. 列出三个”loss 降得很好但模型没用”的 bug 和它们的检测方法。
13. torch.compile 后保存 checkpoint 要注意什么？
14. bf16 和 fp16 在是否需要 GradScaler 上有什么区别？

出关标准

• 自己写出 mini-GPT，通过参数量、初始 loss、因果性三个检查
• 在 Tiny Shakespeare 上训到验证集 BPC ≈ 2.1–2.3，生成有剧本格式
• 完成三模型对照表，结论回应了参数量这个混杂因素
• 画出按位置的 loss 曲线并能解释
• 完成至少 3 个消融，每个都有”预测 → 结果 → 解释”
• 读完 Attention Is All You Need 三遍，能讲清 Table 3 的每一行消融
• 能用 §6 的调试清单排查常见 bug

疑问清单