第 2 讲：PyTorch、einops 与资源核算

对应：Stanford CS336 Spring 2026 Lecture 2（PyTorch (einops), resource accounting）。Spring 2025 的 Lecture 2 内容基本相同，有公开录像。 配套代码：code/lecture02/ 预计阅读时间：3–4 小时（第一遍通读 1.5 小时，第二遍对照代码 2 小时）

0. 本讲导航

0.1 学完本讲，你应该能够

1. 用”餐巾纸计算”（napkin math）回答两类问题：训练一个模型要多久、一块 GPU 能放下多大的模型。
2. 计算任意张量占用的内存；解释张量的存储与步长，区分”视图”和”复制”。
3. 说清 FP32、FP16、BF16、FP8、FP4 的区别；解释混合精度训练为什么这样设计。
4. 使用 einops 的 einsum、rearrange、reduce、repeat 操作张量，并知道它防住了哪一类 bug。
5. 推导矩阵乘法的 FLOPs（$2mkn$），推导”反向 ≈ 2 倍前向”，从而得到训练计算量 ≈ $6ND$。
6. 定义并计算 MFU（模型浮点利用率）；解释为什么逐元素运算”FLOPs 很少，时间却不少”。
7. 估算训练时的显存：参数、梯度、优化器状态、激活值；理解梯度累积和激活检查点的取舍。
8. 从 SGD 一路推演到 AdamW，实现学习率预热 + 余弦退火和梯度裁剪。
9. 写出一个完整的训练循环：数据加载、批次、交叉熵、训练、评估、检查点、生成。

0.2 先修知识

需要安装 PyTorch、numpy、einops（见代码 README）。

0.3 本讲地图

第1节  为什么要做资源核算？（两个动机问题）
   │
   ├── 内存这条线 ──────────────────────┐        计算这条线 ──────────────────┐
   │   第2节 张量：内存 = 元素数 × 字节数    │        第5节 FLOPs：2mkn、6ND、MFU   │
   │   第3节 浮点格式：每个数用几个字节？     │        第6节 算术强度初步             │
   │   第7节 显存：参数/梯度/优化器/激活       │                                     │
   │                                        │                                     │
   第4节 einops：写对张量操作                                                      │
   │                                                                              │
第8节 模型与初始化 → 第9节 优化器 → 第10节 训练循环（把一切组装起来）

本讲是整门课的”会计课”：之后每一讲讨论的架构、并行、推理、规模定律，本质上都是在这里建立的两本账（计算账和内存账）上做优化。

1. 为什么要做资源核算？

1.1 两个动机问题

CS336 用两个问题开场，你读完本讲应该能在几分钟内用纸笔算出答案：

问题一：在 1024 块 H100 上，用 15 万亿（15T）个 token 训练一个 700 亿（70B）参数的模型，要多久？

问题二：用 AdamW 训练，8 块 H100（每块 80 GB）最多能放下多大的模型？

如果每个参数的字节数按不同的训练方案算成 12，就能放下约 530 亿参数。答案取决于你对每一项的精度假设——这正是本讲要讲清楚的。

1.2 “餐巾纸计算”的思维方式

大模型团队在开始任何实验之前，都会先做这种粗略但数量级正确的估算。原因是：

• 实验太贵。一次大规模训练动辄数百万美元，你必须事先知道它大概要多久、要多少卡、会不会爆显存。
• 找瓶颈。算完才知道时间花在哪（计算？内存搬运？通信？），优化才有方向。
• 判断”好不好”。你的训练跑出 MFU 30%，这算好还是差？没有核算就无从判断。

需要核算的资源主要有三种：计算（FLOPs）、内存（字节）、通信（多卡之间传的字节，第 7–8 讲）。本讲处理前两种。

2. 张量与内存

2.1 一切都是张量

深度学习里的所有数据——输入、参数、梯度、优化器状态、中间激活——都是张量（tensor），也就是多维数组。它的维数（rank）就是有几个下标：0 维是标量，1 维是向量，2 维是矩阵。Transformer 里最常见的是 4 维张量，例如注意力中的（批次, 头, 序列位置, 头维度）。

2.2 内存 = 元素个数 × 每个元素的字节数

这是资源核算最基本的公式：

例如 GPT-3 某一层前馈网络的一个权重矩阵形状是 $12288 \times 49152$，用 FP32（4 字节）存储：

GPT-3 有 96 层，每层有两个这样大小的前馈矩阵，光这部分就要约 464 GB（FP32）。一块 H100 只有 80 GB。

一个实用技巧：PyTorch 的 device="meta" 可以创建”只有形状、没有数据”的张量，它不占内存，但可以正常查询大小。用它来做资源核算非常方便（配套代码 01_tensors_and_memory.py 演示 2）。

2.3 存储与步长：张量在内存里是一条长带

内存本身是一维的。一个 $3\times4$ 的矩阵在内存里是 12 个数排成一排（默认按行排：先第 0 行，再第 1 行……）。PyTorch 的张量由两部分组成：

• 存储（storage）：一维的数据长带；
• 元信息：形状（shape）、步长（stride）、起始偏移。

步长告诉你”沿某一维走一步，要在长带上跳过几个元素”。$3\times4$ 矩阵的步长是 $(4, 1)$：沿行方向走一步要跳 4 个，沿列方向走一步跳 1 个。于是：

理解步长的好处是，你会明白哪些操作是”免费”的：

配套代码的实测：对 $4096\times4096$ 的张量，转置（视图）约 0.6 微秒，而”转置并复制”约 13 毫秒——差了四个数量级，因为后者要读写全部 64 MB 数据。

两个常见的坑：

1. 视图共享数据：修改切片会改变原张量（row = x[1]; row[0] = 100 之后 x[1, 0] 也变成了 100）。
2. 转置后不能直接 view：x.T.view(-1) 会报错，因为转置后的张量在内存中不连续。

为什么这对大模型重要？ 显存核算时，视图不额外占显存，但每个运算的输出都会；第 5、6 讲写 GPU 核函数时，你要亲手用步长计算每个元素的地址。

2.4 设备

张量默认在 CPU 内存里。用 x.to("cuda")（NVIDIA GPU）或 x.to("mps")（Apple 芯片）把它复制到加速器上。要点：

• 训练时应该让数据和参数一直留在 GPU 上，避免反复搬运。在 NVIDIA 机器上，CPU 与 GPU 之间通过 PCIe 连接，带宽比 GPU 显存带宽低一两个数量级。
• GPU 上的运算是异步的：Python 发出命令后立刻返回。所以计时前必须”同步”（等 GPU 算完），否则测到的只是”发命令”的时间（第 5 节和第 6 讲）。

3. 浮点数格式：每个数用几个字节？

想省内存、省带宽、算得更快，最直接的办法是每个数用更少的比特。但比特少了，数就表示得不准。理解这里的取舍，是理解”混合精度训练”的关键。

3.1 浮点数的结构

一个浮点数由三部分组成：

• 指数位决定范围：能表示多大和多小的数；
• 尾数位决定精度：相邻两个可表示的数挨得有多近。

总位数固定时，两者此消彼长：

（数据来自配套代码 02_floating_point.py 演示 1 中的 torch.finfo。）

3.2 FP16 的问题：范围太窄

FP16 最大只能表示 65504，最小正规数约 $6\times10^{-5}$。实验：

• $10^{-8}$ 存成 FP16 → 0（下溢）。梯度常常就是这么小的数，用 FP16 训练时大量梯度会”消失”。
• 70000 存成 FP16 → inf（上溢）。激活值偶尔很大时就会出问题。

早期用 FP16 训练需要损失缩放（loss scaling）：先把损失乘以一个大数（比如 1024），让梯度被放大到 FP16 能表示的范围，更新参数前再除回去。这很麻烦，而且缩放因子要动态调整。

3.3 BF16：砍尾数，保指数

BF16（brain float 16，Google Brain 为深度学习设计）的思路是：指数位和 FP32 一样多（8 位），只砍尾数。于是：

• 范围和 FP32 一样：$10^{-8}$ 和 70000 都能正常表示，不需要损失缩放。
• 精度很差：只有 2~3 位有效数字。在 1 附近，相邻两个 BF16 数相差 $2^{-7} \approx 0.0078$，所以 $1 + 0.001 = 1$；在 1024 附近间隔是 8，所以 $1024 + 1 = 1024$。
• 转换极其便宜：BF16 的比特恰好是 FP32 的前 16 位（演示 2），FP32 → BF16 只需要截掉后 16 位（并四舍五入）。

为什么精度这么差还能用？因为神经网络对单个数的噪声很鲁棒。真正危险的是累加：

3.4 累加误差：为什么累加必须用 FP32

实验：把 0.01 逐个加 10000 次（正确答案 100）。

BF16 累加到 4 就”卡住”了：在 $[4, 8)$ 之间，相邻数的间隔是 $1/32 \approx 0.031$，而 0.01 小于间隔的一半，加上去会被四舍五入回原值。同样的道理：

• 参数更新：$\theta \leftarrow \theta - \text{lr}\times\text{更新量}$ 中，更新量往往比 $\theta$ 小几个数量级。如果 $\theta$ 用 BF16 存储，从 1.0 开始减 1000 次 $10^{-4}$，结果纹丝不动（仍是 1.0，正确答案是 0.9）。
• 求和类运算：矩阵乘法中的累加、LayerNorm 的均值方差、softmax 的分母、损失的平均，都是大量数相加。

3.5 混合精度训练

综合以上，混合精度训练的原则是：

PyTorch 的 torch.autocast 按一张规则表自动为运算选精度：矩阵乘法用 BF16，损失函数等提升为 FP32。注意这张表因设备而异（配套代码演示 6：在 CPU 上 softmax 仍是 BF16，在 CUDA 上会提升为 FP32），所以关键的归约最好自己显式用 .float()。

3.6 更低的精度：FP8 与 FP4

H100 起硬件支持 FP8，最新的 B200 支持 FP4。它们粗糙得惊人——FP4（E2M1）连正负号在内一共只能表示 15 个不同的数：$\pm\{0, 0.5, 1, 1.5, 2, 3, 4, 6\}$。怎么还能用？关键是缩放因子：

• 整体缩放：先除以一个缩放因子（通常 = 最大绝对值 / 格式能表示的最大值），把数值挪到格式擅长的范围，量化后再乘回来。实验：一组量级 $10^{-4}$ 的梯度直接转 FP8 E4M3，100% 变成 0；按最大值缩放后，0% 变成 0，相对误差只有 2.7%。
• 分块缩放：大模型的激活里常有少数”离群值”（比其他值大一两个数量级）。整体缩放时，缩放因子被离群值撑大，其余正常值都被挤到 0 附近。实验中，整体缩放的 FP4 让100% 的正常值变成了 0；改为每 16 个元素一个缩放因子后，只有 10% 变成 0，相对误差从 17.5% 降到 2.6%。

这就是 MXFP8（每 32 个元素一个缩放因子）、NVFP4（每 16 个一个）等新格式的核心思想。FP8 训练已经在一些前沿模型中使用；FP4 训练是 2025 年起的前沿方向。

4. einops：用维度名字写张量操作

4.1 动机：最危险的 bug 不会报错

写 Transformer 时，最容易出错的是维度操作。更可怕的是，很多错误不会报错。配套代码 03_einops.py 演示 5：线性层输出形状是 (batch, seq, heads×dim)，想拆成 (batch, heads, seq, dim)。

wrong = x.view(B, H, S, D)                              # 形状对了，内容错了！
right = x.view(B, S, H, D).transpose(1, 2)              # 正确
right = rearrange(x, "b s (h d) -> b h s d", h=H)       # einops 写法，一目了然

x.view(B, H, S, D) 只是把内存里的数按新形状重新读一遍，但内存里的顺序是 (seq, heads, dim)。结果形状完全正确，模型照样能训练，只是效果莫名其妙地差——这种 bug 可能要花几天才能找到。

einops 的思路是给每个维度起名字，用一个字符串描述”输入长什么样 → 输出长什么样”。代码即文档。

4.2 四个核心函数

einsum：广义的矩阵乘法。规则只有一条——在输入中出现、但输出中没有的维度，会被求和：

einsum(x, y, "seq1 hidden, hidden seq2 -> seq1 seq2")          # 矩阵乘法（hidden 被求和）
einsum(q, k, "b h q_pos d, b h k_pos d -> b h q_pos k_pos")    # 注意力分数（d 被求和）
einsum(a, b, "i, j -> i j")                                     # 外积（没有维度被求和）

rearrange：转置、拆分、合并维度。括号表示”这一维是由几个维度合并来的”：

rearrange(x, "batch seq (heads dim) -> batch heads seq dim", heads=8)   # 拆出注意力头
rearrange(img, "b c (h p1) (w p2) -> b (h w) (p1 p2 c)", p1=8, p2=8)   # 图像切块（第 17 讲）

reduce：沿指定维度求和、平均、最大值：reduce(x, "batch seq hidden -> batch hidden", "mean")。

repeat：复制。例如分组查询注意力（GQA，第 3 讲）中，把 2 个 KV 头复制给 8 个查询头：repeat(k, "b kv s d -> b (kv g) s d", g=4)。

4.3 形状注释

CS336 推荐用 jaxtyping 库给函数参数写上形状，例如 q: Float[Tensor, "batch heads seq d"]。即使不做运行时检查，它们也是极好的文档。本课程的代码在文档字符串的【形状】里用同样的风格描述。

5. 计算量核算：FLOPs、6ND 与 MFU

5.1 两个容易混淆的词

• FLOPs（小写 s）：浮点运算的次数，衡量”要做多少工作”。
• FLOP/s（有时写作 FLOPS）：每秒的浮点运算次数，衡量”硬件有多快”。

几个参考数字：GPT-3 的训练约 $3.14\times10^{23}$ FLOPs；据推测 GPT-4 约 $2\times10^{25}$ FLOPs。硬件方面（BF16 稠密算力）：

注意 BF16 算力是 FP32 的十几倍——这是 Tensor Core 专门为低精度矩阵乘法设计的结果，也是混合精度训练的另一个动机。（厂商宣传的数字常常是”带结构化稀疏”的，比稠密值高一倍，比如 H100 的 1979 TFLOP/s；训练中基本用不上稀疏加成，核算时要用稠密值。）

5.2 矩阵乘法：$2mkn$

$(m\times k)$ 乘 $(k\times n)$，结果有 $mn$ 个元素，每个是 $k$ 对数的乘积之和（$k$ 次乘法 + $k$ 次加法），所以：

配套代码 04_flops_and_mfu.py 用 PyTorch 自带的 FLOP 计数器验证了这个公式。对一个线性层 $y = xW$（$x$: $B\times D_{in}$，$W$: $D_{in}\times D_{out}$），前向 FLOPs $= 2 \times B \times D_{in}D_{out} = 2 \times (\text{token 数}) \times (\text{这层的参数量})$。

逐元素运算（加法、ReLU 等）的 FLOPs 只和元素个数成正比，比矩阵乘法少几个数量级，计算量核算时通常忽略（但它们的时间不能忽略，见第 6 节）。

5.3 反向传播：为什么是前向的 2 倍

对一层 $h_{out} = h_{in} W$，前向是一次矩阵乘法（$2BD_{in}D_{out}$）。反向传播需要两样东西：

1. 对参数的梯度 $\dfrac{\partial L}{\partial W} = h_{in}^\top \dfrac{\partial L}{\partial h_{out}}$，用来更新这一层 → 一次同样大小的矩阵乘法；
2. 对输入的梯度 $\dfrac{\partial L}{\partial h_{in}} = \dfrac{\partial L}{\partial h_{out}} W^\top$，用来继续往前传给上一层 → 又一次。

所以反向 = 2 × 前向，前向 + 反向 = 3 × 前向：

实验验证（配套代码演示 2、3）：

• 两层网络，输入也需要梯度时，反向/前向 = 2.00；输入是数据（不需要梯度）时 = 1.50——第一层省掉了”对输入的梯度”那次矩阵乘法。层数越多，这个省掉的比例越小。
• 8 层网络，实测总 FLOPs / (N × D) = 5.75，接近 6。

6ND 没有算进什么？

• 注意力中与序列长度相关的计算（第 3 讲）：序列很长时不可忽略；
• 激活重计算（第 7.4 节）：多做一次前向，变成约 $8ND$；
• 逐元素运算：FLOPs 可以忽略，时间不能。

5.4 MFU：模型浮点利用率

例如：在 1 块 H100 上训练 15 亿参数的模型，实测 40,000 tokens/秒，则 MFU $= 6 \times 1.5\times10^9 \times 40000 / 989.5\times10^{12} \approx 36\%$。MFU 达到 50% 左右就算相当好了。

MFU 只把”模型必需的”计算（每个 token $6N$）算作有效工作，激活重计算之类的额外计算不算——所以它衡量的是硬件被”有效利用”的程度。

5.5 实测：为什么大矩阵才能跑满硬件

在一台 Apple M4 Pro 笔记本上实测 FP32 方阵乘法（你的电脑数字会不同，但规律一样）：

小矩阵远远跑不满：64×64 时 GPU 只达到峰值的 0.2%。原因是计算太少，启动运算、调度线程等固定开销占了大头，而且相对计算量而言数据搬运太多。这就是训练要用大批次、大矩阵的原因。

6. 算术强度初步：FLOPs 不等于时间

同一台电脑上，对 1677 万个 FP32 数做 x * 2：只达到约 26 GFLOP/s，而矩阵乘法能达到约 3,300 GFLOP/s，差了一百多倍。

原因在于算术强度（arithmetic intensity）：

• x * 2：每个元素读 4 字节、写 4 字节，只做 1 次乘法，强度 = 1/8 FLOP/字节。速度完全被”搬数据”卡住（内存受限）。
• $n\times n$ 矩阵乘法：做 $2n^3$ 次运算，只需读写约 $3n^2$ 个数，强度随 $n$ 增长。每个数据被重复使用很多次（计算受限）。

硬件的”平衡点”= 峰值算力 / 内存带宽。H100 是 $989.5 \times 10^{12} / 3.35\times10^{12} \approx 295$ FLOP/字节：算术强度低于它的运算都是内存受限的。Transformer 里的 LayerNorm、激活函数、残差相加、softmax 都是这类运算——FLOPs 占比不到 1%，却可能占训练时间的相当一部分。第 5 讲会用 Roofline 模型系统地分析这件事，第 6 讲用算子融合解决它。

7. 显存核算

训练时显存里要放四样东西：

7.1 每个参数几个字节？

最常用的是第三种：前向和反向用 BF16 副本计算（快、省带宽），优化器用 FP32 主权重更新。按 16 字节/参数：70 亿参数的模型光这三项就要 112 GB，一块 80 GB 的 GPU 放不下；700 亿参数要 1120 GB，至少 14 块卡。第 7 讲的 ZeRO / FSDP 会把这 16 字节切分到多块 GPU 上。

配套代码 05_memory_accounting.py 演示 2 实测了 PyTorch 的 AdamW：step() 之后，优化器状态恰好是参数字节数的 2 倍（exp_avg 和 exp_avg_sq 两个与参数同形状的 FP32 张量）。注意状态是懒创建的：第一次 step() 之前不占内存。

7.2 激活值：精确测量

前向传播时，PyTorch 会把反向传播需要的中间结果”存起来”，这就是激活值。配套代码用 torch.autograd.graph.saved_tensors_hooks 拦截每一次”保存”，精确测量了一个 $L$ 层（线性层 + ReLU）网络保存的激活：

完全吻合：每层保存一个 $B\times D$ 的张量（ReLU 的输出，它同时是下一层的输入），外加最初的输入。激活值与层数、批次大小成正比，参数量与批次无关——批次足够大时，激活会超过参数成为显存的大头。Transformer 每层保存的激活远不止一个 $B\times D$，而且与序列长度有关（第 3 讲给出公式）。

7.3 梯度累积：用时间换显存

想用大批次（比如 64）训练，但显存只放得下 16？把大批次拆成 4 个小批次，依次做前向和反向，梯度会自动累加到 .grad 中，4 次之后再更新参数。实验（演示 4）：

• 两种方式的梯度最大差异只有 $2.3\times10^{-10}$（浮点舍入误差），数学上完全等价；
• 激活显存降到 1/4。

两个要点：每个小批次的损失要除以小批次个数（否则梯度大 4 倍）；代价是矩阵变小，硬件利用率下降。

PyTorch 小知识：多次调用 backward() 时梯度是累加的，不会覆盖。这既是梯度累积的实现基础，也是”忘了 zero_grad()”这个经典 bug 的来源。

7.4 激活检查点：用计算换显存

激活检查点（activation checkpointing，也叫重计算）的思路：前向时不保存某些激活，反向时需要它们再重新算一遍。

如果每 $k$ 层设一个检查点：前向只保存每段的输入（$L/k$ 份），反向时逐段重算，临时保留一段内的 $k$ 份激活。峰值显存 $\approx L/k + k$，在 $k=\sqrt{L}$ 时最小，约 $2\sqrt{L}$。$L=100$ 时，激活显存从 100 份降到约 20 份。

代价可以用 FLOP 计数器精确验证（演示 5）：不做检查点时，前向+反向 = 2.94 倍前向；每层都做检查点时 = 3.94 倍——恰好多了一次前向。训练计算量从约 $6ND$ 变成约 $8ND$。实践中常常只对”激活大、重算便宜”的部分做检查点（选择性重计算）。

8. 模型：参数与初始化

8.1 参数在 PyTorch 里是什么

模型的参数是 nn.Parameter——一种会被自动登记、能求梯度的张量。自定义一个层时，继承 nn.Module，在 __init__ 里把 nn.Parameter 赋给属性，在 forward 里写计算。CS336 作业 1 要求你自己实现 Linear 和 Embedding，而不是用 nn.Linear，配套代码 07_training_loop.py 给出了示例。

8.2 为什么初始化的尺度如此重要

考虑 $y = xW$，$x$ 的每个元素方差为 1，$W$ 的每个元素方差为 $\sigma^2$，彼此独立。$y$ 的每个元素是 $d_{in}$ 个乘积之和，所以：

想让信号的尺度在层与层之间保持不变，就要 $\sigma = 1/\sqrt{d_{in}}$。实验（$d=512$，10 个线性层）：

信号爆炸或消失时，梯度同样会爆炸或消失，训练根本无法进行。CS336 作业 1 用的是最后一种：Xavier 初始化的一种形式（$d_{in}=d_{out}$ 时与 $1/\sqrt{d_{in}}$ 相同），并把超过 $\pm3\sigma$ 的值截断，避免极端值。嵌入层用标准差为 1 的截断正态（因为它是查表，没有求和）。

9. 优化器：从 SGD 到 AdamW

优化器决定”拿到梯度后怎么更新参数”。它影响收敛速度和稳定性，也直接决定显存（每个参数要额外存几个数）。沿着历史脉络看，每一步都在解决上一步的问题：

9.1 为什么需要”每个参数自己的学习率”

配套代码 06_optimizers.py 演示 3 在一个”又窄又长的山谷”上比较：$L(w) = \frac12(100 w_1^2 + w_2^2)$，$w_1$ 方向的曲率是 $w_2$ 的 100 倍（起点损失 50.5）：

SGD 的学习率被陡峭方向卡死（必须 $< 2/100$，0.021 就发散），平缓方向每步只缩小约 2%。自适应方法用梯度大小给每个方向”归一化”，两个方向步长相当，能用大得多的学习率。（Adam 前 10 步冲过了头，是动量的”惯性”所致，之后很快收敛。）这个例子的曲率恰好沿坐标轴，所以”每坐标一个学习率”完美对症；真实网络没这么理想，但神经网络不同方向的曲率确实可以相差几个数量级，这是 Adam 家族一统天下的主要原因。

9.2 AdamW 的完整算法

CS336 作业 1 的写法（Loshchilov & Hutter 2019）：

为什么要偏差修正？ $m$、$v$ 都从 0 开始，前几步会被”拉向 0”：第 1 步 $m = (1-\beta_1)g = 0.1g$。除以 $1-\beta_1^t$ 正好修正回来（$t=1$ 时 $0.1g/0.1 = g$）。随着 $t$ 增大，$\beta^t\to0$，修正自动失效。一个直观的推论：Adam 第一步的更新量约为 $\eta\cdot\text{sign}(g)$，与梯度大小无关。

为什么叫”解耦的”权重衰减？ 早期做法是把 L2 正则 $\lambda\theta$ 加进梯度，但这样它也会被 $\sqrt{v}$ 缩放：梯度大的参数衰减少，梯度小的衰减多。AdamW 把权重衰减从梯度中拿出来，直接作用于参数。

一个值得知道的细节：配套代码实现了两个版本。按 PyTorch 官方的计算顺序（先衰减、$\epsilon$ 加在”除以 $\sqrt{1-\beta_2^t}$ 之后”），我们的实现与 torch.optim.AdamW 逐比特一致。按作业 1 的公式，第 1 步参数就有 $8.7\times10^{-4}$ 的差异，50 步后增大到 $6\times10^{-2}$。原因是：两者只差在 $\epsilon$ 的位置，这对梯度接近 0 的参数（比如几乎不激活的 ReLU 神经元）会显著改变更新量；而 Adam 的更新量约等于 $\eta\cdot\text{sign}(g)$，微小的差异就能让参数朝不同方向走一步，之后越滚越大。这不是 bug，两种写法都合理，训练效果也没有区别——但它提醒我们，同样叫”AdamW”，不同框架的实现细节可能不同。

9.3 学习率调度：预热 + 余弦退火

• 为什么预热：刚开始参数是随机的，梯度方向很不稳定，Adam 的 $v$ 也还没估计准；一上来就用大学习率容易”一步跳飞”。
• 为什么退火：训练后期需要小步长才能落进更低的谷底，而不是在谷底附近来回跳。

第 11 讲会介绍另一种流行的调度 WSD（预热–稳定–衰减）。

9.4 梯度裁剪

把所有参数的梯度拼成一个大向量，算它的 L2 范数 $\|g\|$；如果超过阈值 $M$，就把所有梯度乘以 $M/(\|g\|+\epsilon)$。方向不变，只缩短长度。它防止偶尔出现的”梯度爆炸”一步毁掉之前的训练。注意必须用全局范数，而不是每个参数分别裁剪（否则会改变梯度方向）。配套代码的实现与 torch.nn.utils.clip_grad_norm_ 的差异在 $10^{-8}$ 量级。

10. 训练循环：把一切组装起来

配套代码 07_training_loop.py 实现了一个完整的训练程序，包含 CS336 作业 1 训练部分要求的全部要素。模型刻意用了一个简单的”固定窗口”神经语言模型（Bengio 等人 2003 年的思路：用前 $k$ 个 token 的嵌入拼接后经过 MLP 预测下一个），因为本讲的重点是训练的”基础设施”。第 3 讲把模型换成 Transformer，其余部分几乎不用改。

10.1 数据：uint16 文件 + memmap

真实训练中，先用分词器把整个语料编码一次，把 token id 存成一个大的二进制文件，训练时直接读取：

• 为什么 uint16：作业的词表是 1 万～3.2 万，不超过 65536，每个 token 2 字节就够了，比 int64 省 4 倍。
• 为什么 memmap：np.memmap 不把整个文件读进内存，而是让操作系统按需加载。即使文件有 100 GB，打开也几乎不花时间和内存；每次取批次只读几十 KB。
• 坑：二进制文件不记录数据类型。用 uint8 读 uint16 文件，每个 token 会被拆成两半，而且不会报错。

10.2 取批次

随机选若干个起点 $i$，输入 $x = \text{data}[i : i+L]$，目标 $y = \text{data}[i+1 : i+L+1]$（右移一位）。一个长度为 $L$ 的样本同时提供了 $L$ 个”预测下一个 token”的任务。起点的上界是 len(data) - L - 1，否则 $y$ 会越界。

10.3 数值稳定的交叉熵

减去最大值后所有指数都 $\le 0$，不会上溢；至少有一项等于 1，求和不会下溢。实验：logits 乘以 1000 后，朴素实现得到 nan，稳定实现正常。另一个有用的参考点：刚初始化的模型损失应接近 $\ln V$（均匀猜测），字节级模型是 $\ln 256 \approx 5.55$。

10.4 训练、评估与过拟合

每一步：设置学习率 → 取批次 → 前向算损失 → zero_grad → backward → 梯度裁剪 → step。每隔一段时间在验证集上评估（用 torch.no_grad()，不记录计算图）。

在第 1 讲的小语料上训练 1000 步的结果：

• 比特/字节 = 损失 / $\ln 2$。第 1 讲说过”预测就是压缩”：这就是用该模型压缩文本时每个字节平均需要的比特数（原始 UTF-8 是 8）。
• 验证损失在第 200 步最低，之后训练损失继续下降、验证损失却回升——过拟合。原因是训练集只有约 1.6 万字节，1000 步 × 2048 token 相当于把它重复看了约 129 遍，模型开始”背诵”。真实大模型的数据极多，每个 token 通常只看一遍，几乎不会遇到这个问题。应对办法：保存验证损失最低的检查点（早停）、更小的模型、更多的数据。

10.5 可复现性：固定随机种子还不够

实验：同样的随机种子、同样的代码训练 50 步，两次的参数并不逐比特相同。原因是：

1. 浮点加法不满足结合律：$(0.1+0.2)+0.3 = 0.6000000000000001$，而 $0.1+(0.2+0.3) = 0.6$。
2. 为了快，很多运算用多个线程并行求和（比如嵌入层的反向要把同一个 token 在不同位置的梯度加到同一行），线程完成的先后每次不同，加法顺序就不同。

torch.use_deterministic_algorithms(True) 强制使用确定性实现后，两次结果完全一致。大规模训练通常不追求逐比特可复现（太慢），但调试时，能复现是定位 bug 的前提。

10.6 检查点

保存检查点要存：模型参数、优化器状态、当前步数、随机数生成器状态。实验：训练 100 步 → 保存 → 用全新的对象恢复 → 再训练 100 步，与一口气训练 200 步的结果逐比特一致。如果只恢复模型、不恢复优化器，最大参数差异达到 0.088——AdamW 的 $m$、$v$ 被清零，相当于训练中途又”重启”了一次优化器。

10.7 生成与温度

生成时重复：算出下一个 token 的 logits → 除以温度 $T$ → softmax → 按概率采样 → 接到末尾。$T<1$ 让分布更尖（更保守），$T>1$ 让分布更平（更随机），$T\to0$ 就是贪心解码。这个只能看到最近 8 个字节的模型，能拼出一些真实的单词（and、the、said、Mom），也会造出很多似是而非的伪单词，句子毫无逻辑。第 3 讲的 Transformer 可以看到整个上下文。

11. 本讲小结

1. 资源核算是大模型工程的基本功：先用餐巾纸计算估算时间和显存，再动手。
2. 内存 = 元素个数 × 每个元素的字节数。视图不复制数据，多数运算的输出会。
3. BF16 用 FP32 的范围换掉了精度；累加、归约、主权重、优化器状态必须用 FP32——这就是混合精度训练。FP8/FP4 依靠（分块）缩放因子才能使用。
4. einops 用维度名字写张量操作，防住”形状对、内容错”的隐蔽 bug。
5. 矩阵乘法 $2mkn$；反向 ≈ 2 倍前向；训练 ≈ $6ND$ FLOPs。MFU = 实际模型 FLOP/s ÷ 峰值，50% 左右算好。
6. FLOPs 不等于时间：逐元素运算内存受限，小矩阵跑不满硬件。
7. 显存 = 参数 + 梯度 + 优化器状态（混合精度 AdamW 共约 16 字节/参数）+ 激活值（随批次和层数增长）。梯度累积和激活检查点分别用时间、用计算换显存。
8. 初始化要让方差保持稳定（$\sigma\approx1/\sqrt{d_{in}}$）。
9. AdamW = 动量 + 自适应学习率 + 偏差修正 + 解耦权重衰减，代价是每参数 2 个额外状态。配合预热 + 余弦退火、梯度裁剪使用。
10. 训练循环的要素：memmap 数据、随机批次、稳定交叉熵、评估、早停、确定性、完整的检查点、温度采样。

12. 与作业 1 的关系

作业 1 中与本讲直接对应的部分：

作业的资源核算题（例如”GPT-2 XL 用 AdamW 训练需要多少显存”）需要先知道 Transformer 的结构，第 3 讲会完整地算一遍。

13. 练习

练习 1（内存）：一个形状为 $(32, 1024, 4096)$ 的 BF16 张量占多少内存？

$32\times1024\times4096\times2 = 268{,}435{,}456$ 字节 $\approx 268$ MB（256 MiB）。这正是”批次 32、序列 1024、隐藏维 4096”的一份激活值的大小。

练习 2（步长）：一个连续的 $(2, 3, 4)$ 张量的步长是多少？元素 $[1, 2, 3]$ 在存储中的位置是多少？对它做 transpose(0, 2) 后，形状和步长是多少？

步长 $(12, 4, 1)$；位置 $1\times12 + 2\times4 + 3\times1 = 23$（也就是最后一个元素）。transpose(0, 2) 后形状 $(4, 3, 2)$，步长 $(1, 4, 12)$，不连续。

练习 3（BF16）：在 BF16 中，$1024 + 1$、$1024 + 4$、$1024 + 5$ 分别等于多少？为什么？

BF16 有 7 位尾数（加上隐含的 1 共 8 位有效二进制位），在 $[1024, 2048)$ 之间相邻数的间隔是 $2^{10-7} = 8$。所以 $1025 \to 1024$；$1028$ 恰好在 1024 和 1032 中间，按”四舍五入到偶数”规则取 1024；$1029 \to 1032$。可以用 torch.tensor(1024., dtype=torch.bfloat16) + 5 验证。

练习 4（FLOPs）：$(4096\times8192) @ (8192\times1024)$ 需要多少 FLOPs？在 H100 上以峰值 BF16 算力计算需要多久？

$2\times4096\times8192\times1024 \approx 6.87\times10^{10}$ FLOPs；$6.87\times10^{10} / 989.5\times10^{12} \approx 69$ 微秒。

练习 5（训练时间）★：用 8 块 H100、MFU 40%，训练一个 10 亿参数的模型、200 亿 token（约为 Chinchilla 最优的 20 token/参数），需要多久？

$6\times10^9\times2\times10^{10} = 1.2\times10^{20}$ FLOPs；有效算力 $8\times989.5\times10^{12}\times0.4 \approx 3.17\times10^{15}$ FLOP/s；时间 $\approx 3.8\times10^4$ 秒 $\approx 10.5$ 小时。

练习 6（反向传播）：一个 3 层线性网络（每层 $D\times D$，批次 $B$），输入是数据（不需要梯度），前向和反向的 FLOPs 分别是多少？比值是多少？

前向 $3\times2BD^2 = 6BD^2$。反向：3 次”对参数的梯度”+ 2 次”对输入的梯度”（第一层不需要）$= 5\times2BD^2 = 10BD^2$。比值 $10/6 \approx 1.67$。层数为 $L$ 时比值为 $(2L-1)/L$，$L$ 越大越接近 2。

练习 7（Adam 的第一步）：$\beta_1 = 0.9, \beta_2 = 0.999$，某参数第 1 步的梯度 $g = 2$，求 $m$、$v$、偏差修正后的 $\hat m = m/(1-\beta_1)$、$\hat v = v/(1-\beta_2)$，以及更新量（忽略 $\epsilon$ 和权重衰减）。

$m = 0.1\times2 = 0.2$，$v = 0.001\times4 = 0.004$；$\hat m = 2$，$\hat v = 4$；更新量 $= \eta\times\hat m/\sqrt{\hat v} = \eta\times2/2 = \eta$。第一步的更新量恰好是 $\eta\cdot\text{sign}(g)$，与梯度大小无关。

练习 8（激活检查点）：一个 64 层的网络，最优的检查点间隔是多少？峰值激活显存（以一层为单位）从多少降到多少？代价是什么？

间隔 $\sqrt{64} = 8$；峰值从 64 降到 $64/8 + 8 = 16$。代价是多做一次前向，训练计算量从约 $6ND$ 增加到约 $8ND$。

练习 9（思考）：(a) 如果在训练循环中忘了 optimizer.zero_grad()，会发生什么？(b) 为什么评估时要用 torch.no_grad()？

(a) 梯度会在每一步不断累加，相当于用”从开始到现在所有梯度之和”来更新，更新量越来越大，训练很快发散。(b) 评估不需要反向传播，不记录计算图可以省下保存激活的显存，也更快。

练习 10（动手）： (a) 在 07_training_loop.py 中把 window 从 8 改成 16、32，最佳验证损失如何变化？过拟合是更早还是更晚出现？ (b) 实现 top-k 采样：只在概率最高的 $k$ 个 token 中采样。 (c) 用 06_optimizers.py 中我们自己写的 AdamW 替换 07_training_loop.py 中的 torch.optim.AdamW，训练结果有什么变化？

14. 延伸阅读

• CS336 Spring 2026 / 2025 Lecture 2 的可执行讲义（GitHub stanford-cs336/lectures 中的 lecture_02.py）。
• einops 官方教程：einops.rocks，其中的”Basics”部分约 30 分钟。
• Micikevicius et al.，Mixed Precision Training（2018）：混合精度训练的经典论文，损失缩放的来源。
• NVIDIA，Training with Mixed Precision 文档，以及 FP8 / MXFP8 / NVFP4 相关的技术博客。
• Kingma & Ba，Adam: A Method for Stochastic Optimization（2014）；Loshchilov & Hutter，Decoupled Weight Decay Regularization（2019）。
• Chen et al.，Training Deep Nets with Sublinear Memory Cost（2016）：激活检查点的 $\sqrt{L}$ 方案。
• Horace He，Making Deep Learning Go Brrrr From First Principles（博客）：计算受限、内存受限、开销受限的直观解释，第 5 讲的预习材料。

15. 下一讲预告

本讲搭好了训练的”基础设施”，但模型还只是一个只能看 8 个字节的小 MLP。下一讲（架构与超参数）把它换成现代 Transformer：

• 从 2017 年的原始 Transformer 到今天的 Llama / Qwen，改了哪些地方？为什么？
• Pre-norm 还是 Post-norm？LayerNorm 还是 RMSNorm？为什么 FLOPs 很少的归一化层也值得优化？
• SwiGLU 为什么是现在的主流？前馈层的宽度为什么是 $\frac83 d$？
• RoPE 如何用”旋转”来编码相对位置？
• 头的数量、深宽比、词表大小……大家的超参数为什么惊人地一致？
• 训练不稳定时，z-loss、QK-norm、logit 软截断分别在防什么？

下一讲之前，建议你：跑通 code/lecture02/ 下的 7 个脚本，完成练习 5、6、7，并尝试不看代码自己写出 AdamW 和 get_batch。