PyTorch 与资源核算

0. 本讲导读

上一讲讲了课程总览和分词（tokenization），把文本变成了整数序列。从这一讲开始，我们真正进入"训练一个语言模型"的工程内核。

这门课反复问的核心问题只有一个：

这就是"资源核算"（resource accounting）。它听上去平淡，但它是这门课全部内容的地基：后面讲架构（Lecture 3）、并行（Lecture 6-7）、推理（Lecture 10）、Scaling Law（Lecture 8），每一处的设计取舍最终都会落回到"这么做省了多少 FLOPs / 省了多少字节 / 打满了多少算力"。不会算账，就只能凭感觉调参。

Percy 在课上明确说了，这一讲希望你带走三样东西：

Percy 特别强调：今天没有任何"机器学习魔法"。不讲 loss 怎么降、不讲哪个 trick 有效，全部是"多少字节、多少次乘加、多少秒"。恰恰因为如此，这一讲的内容是最不容易过时的。

另外一则课程消息：Marin 项目的 1e23 FLOPs 训练跑完了，而且实测结果命中了事先用 scaling law 做出的预测。这正是本课的价值主张——大模型训练不是碰运气，是可以被预测、被核算的工程。

本讲的地图

整讲按"先算显存、再算算力、最后把两者合起来算训练"的顺序推进：

• 显存核算：张量基础 → 浮点数类型 → 内存布局（view）→ GPU 显存
• 算力核算：einops → 矩阵乘法 FLOPs → MFU → 算术强度与 roofline
• 训练全流程核算：深度网络 → 反向传播 FLOPs（$6ND$）→ 优化器状态 → 完整训练循环
• 省显存的手段：梯度累积、激活重计算

1. 先做两道"餐巾纸算术"

在讲任何机制之前，Percy 先扔出两个问题。这两个问题的答案本身不重要，重要的是它们能被算出来，而且只需要一张餐巾纸。

问题一：70B 模型、15T tokens、1024 张 H100，要训多久？

用后面会推导的公式 $\text{FLOPs} = 6 \cdot N \cdot D$（$N$ = 参数量，$D$ = 训练 token 数）：

再算硬件一天能提供多少 FLOPs。H100 的规格书上写着 bf16 峰值 1979 teraFLOP/s，但那个数字是带稀疏性（with sparsity）的，稠密计算只有一半：

实际训练不可能打满峰值，取一个乐观但现实的 MFU = 0.5：

大约 4.8 个月。这个数量级和公开报道的前沿模型预训练周期是吻合的——这也是为什么一次预训练"跑砸了"代价如此之高，为什么这门课要花整整一讲讲 scaling law（用小模型预测大模型）。

问题二：8 张 H100 上用 AdamW 最大能训多大的模型？

这次算的是显存。每个参数在训练时要占多少字节？逐项列：

8 张 H100 每张 80 GB：

约 53B 参数。

Percy 的评价：这只是粗糙的信封背面估算（back-of-the-envelope calculation），但它给了你那种"快速摸清资源量级"的味道。这门课希望你听到任何一个模型配置，都能在三十秒内说出它大概要多少卡、多少天。

2. 张量：一切的载体

张量（tensor）是存储一切东西的基本单元：

• 数据（data）：token id 序列
• 参数（parameters）：模型权重
• 梯度（gradients）：反向传播得到的导数
• 优化器状态（optimizer state）：动量、二阶矩
• 激活值（activations）：前向过程中的中间结果

这五类东西加起来就是你显存里的全部内容。之后每次说"显存不够"，都可以拆成这五项分别追问。

一个真实的例子：DeepSeek v3.2 的权重在 Hugging Face 上的 safetensors index 里就是一张张量清单——每个张量的名字、形状、dtype 全在里面。想知道一个开源模型的结构，读它的权重清单比读论文还直接。

秩（rank）

张量的秩就是维度个数：

import torch

x = torch.zeros(4)        # rank 1（向量）
x = torch.zeros(4, 8)     # rank 2（矩阵）
x = torch.zeros(4, 8, 2)  # rank 3

在 Transformer 里最常见的是 rank 4 张量，因为注意力要把 hidden 维拆成多个头：

B = 32   # Batch size（批大小）
S = 16   # Sequence length（序列长度）
H = 16   # Number of heads（注意力头数）
D = 64   # Hidden dimension per head（每个头的维度）
x = torch.zeros(B, S, H, D)

四个整数 (32, 16, 16, 64) 里没有任何一个能告诉你哪个是 batch、哪个是 head。这正是后面 einops 要解决的问题。

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

张量占多少内存，只由两件事决定：(i) 有多少个值，(ii) 每个值是什么类型。

def get_memory_usage(x: torch.Tensor) -> int:
    return x.numel() * x.element_size()

x = torch.zeros(4, 8)
assert x.dtype == torch.float32   # 默认类型
assert x.numel() == 4 * 8         # 32 个元素
assert x.element_size() == 4      # float32 = 4 字节
assert get_memory_usage(x) == 4 * 8 * 4   # 128 字节

换成真实规模——GPT-3 前馈层里的一个矩阵（隐藏维 12288，FFN 放大 4 倍）：

get_memory_usage(torch.empty(12288 * 4, 12288))
# = 49152 * 12288 * 4 字节
# = 2,415,919,104 字节 ≈ 2.4 GB

一层里的一个矩阵就 2.4 GB（fp32）。GPT-3 有 96 层，每层还有 attention 的四个投影矩阵。这一下就把"为什么大模型要精打细算"讲清楚了。

3. 浮点数类型：从 fp32 一路降到 fp4

张量的元素通常是浮点数。浮点数的位被分成三段：符号位（sign）、指数位（exponent）、尾数位（mantissa / significand）。一个浮点数的值大致是

关键在于这两段位各自控制什么：

fp32（float32 / 单精度）

在传统科学计算里，fp32 是基线，某些场景还要上双精度 fp64。但在深度学习里你可以邋遢得多——梯度下降本身就是带噪声的过程，多一点数值噪声并不致命。这个观察是整个低精度训练的哲学起点。

fp16（float16 / 半精度）

x = torch.zeros(4, 8, dtype=torch.float16)
assert x.element_size() == 2      # 内存减半

x = torch.tensor([1e-8], dtype=torch.float16)
assert x == 0                     # 下溢（underflow）！

$10^{-8}$ 直接变成 0。fp16 能表示的最小非零值（含次正规数）约 $6\times10^{-8}$，比 $10^{-8}$ 还大。而梯度恰恰是最容易变得很小的东西：训练后期、深层网络、小学习率，梯度轻松掉到 $10^{-7}$ 以下。一旦梯度被截成 0，这个参数就再也不更新了——训练表现为"莫名其妙地不收敛"或"loss 突然发散"。

历史上 fp16 训练要靠 loss scaling（把 loss 乘一个大常数，让梯度整体抬进 fp16 的可表示区间，更新前再除回去）才能稳住。这是个能用但很烦的补丁。

bf16（bfloat16 / brain float）

bf16 的设计思路极其干脆：既然深度学习在乎范围不在乎精度，那就把 fp32 的指数位原封不动搬过来，尾数砍掉。

x = torch.tensor([1e-8], dtype=torch.bfloat16)
assert x != 0     # 不下溢！

结果是：bf16 用和 fp16 一样的内存，却有和 fp32 一样的动态范围。唯一的代价是分辨率变差，而这对深度学习影响小得多。另一个实际好处是 fp32 ↔ bf16 的转换非常廉价（直接砍掉低 16 位即可），而且通常不需要 loss scaling。

混合精度训练（mixed precision training）

把上面的事实合起来，训练的处境是：

• 全 fp32 训练：能跑，但显存要 2 倍、算力慢 10 倍以上（Tensor Core 对低精度有专门通路）。
• 全 fp16 甚至全 bf16 训练：有风险，会不稳定。问题主要出在需要"长期累加小量"的地方——优化器的动量是把成千上万步的梯度做指数平均，用 bf16 的 7 位尾数去累加，小增量会被直接吃掉（$1.0 + 0.001$ 在 bf16 里还是 $1.0$）。

解法就是混合精度（Micikevicius et al., 2018）：

PyTorch 提供了自动混合精度（AMP, automatic mixed precision）库，它会在安全的地方自动把计算转成 bf16——比如矩阵乘法可以转，而 exp、softmax 的归约、LayerNorm 的方差累加这类对精度敏感的操作会保持 fp32：

with torch.amp.autocast("cuda", dtype=torch.bfloat16):
    y = x @ w        # 这个 matmul 会在 bf16 下执行
    z = torch.exp(y) # 这类操作 autocast 会保持 fp32

fp8

2022 年，为了机器学习负载专门标准化了 fp8（FP8 Formats for Deep Learning；NVIDIA 也有一篇很好的 FP8 primer）。H100 支持两种变体：

• E4M3：4 位指数、3 位尾数，范围 $[-448, 448]$。精度稍好，用于前向的权重和激活。
• E5M2：5 位指数、2 位尾数，范围 $[-57344, 57344]$。范围更大，用于梯度（梯度的动态范围更宽）。

注意这里的取舍逻辑和 fp16 vs bf16 一模一样：同样的比特数，你只能在范围和精度之间分配。fp8 之所以要设计成两种，是因为前向和反向对这两者的需求刚好相反。

fp4 / nvfp4

2025 年 NVIDIA 推出了 nvfp4：每个值只有 4 比特。它能表示的全部取值就 16 个：

-6, -4, -3, -2, -1.5, -1.0, -0.5, 0.0, 0.5, 1.0, 1.5, 2, 3, 4, 6

看起来完全不可用——但关键在于 nvfp4 为每一小块（block）单独存一个缩放因子（scale factor）。于是整体动态范围其实很宽，只是同一个块内部的值不能相差太远。这是个非常聪明的取舍：神经网络里的张量本来就是局部同量级的。

NVIDIA 的 Nemotron 3 Super 就是用 NVFP4 训练的，说明 4 比特训练已经从论文走进了真实的大规模训练。

4. 张量是"视图"：storage、stride 与 contiguous

上一节把"一个张量占多少字节"讲清楚了。但还有一个同样重要的问题：什么时候会产生新的字节？ 这决定了你的代码是在悄悄地复制几 GB 数据，还是零开销。

张量 = 存储（storage）+ 元信息

PyTorch 的张量其实是一层薄薄的"视图"（view），它由两部分组成：

• storage：一段连续的一维内存缓冲区，存着实际的数值。
• 元信息：shape（各维大小）、stride（各维步长）、storage_offset（起始偏移）、dtype、device。

取元素 x[i, j] 时，PyTorch 算的是：

所以只要改改 stride 和 offset，就能在不动一个字节的前提下得到一个"新张量"。这就是零拷贝的来源。

x = torch.arange(12).view(3, 4)
x.stride()          # (4, 1)：行内相邻元素隔 1 个，换一行跳 4 个
x.t().stride()      # (1, 4)：转置只是把 stride 反过来，没有复制
x.data_ptr() == x.t().data_ptr()   # True，共享同一段 storage

哪些操作是零拷贝的？

x = torch.zeros(2, 3)
y = x[0]        # view
y[0] = 1.0
# x[0][0] 现在也变成了 1.0！

5. 从 CPU 到 GPU

默认情况下张量存在 CPU 内存里：

x = torch.zeros(32, 32)
assert x.device == torch.device("cpu")

要用上 GPU 的大规模并行，得把张量搬过去：

device = "cuda" if torch.cuda.is_available() else "cpu"

# 方式一：先在 CPU 上创建再搬过去（会经历一次 PCIe 传输）
x = torch.zeros(32, 32)
x = x.to(device)

# 方式二：直接在 GPU 上创建（推荐，省掉一次传输）
with torch.device(device):
    x = torch.zeros(32, 32)
    assert x.device.type == "cuda"

查看当前占了多少显存：

torch.cuda.memory_allocated()      # 当前已分配（字节）
torch.cuda.max_memory_allocated()  # 峰值
torch.cuda.reset_peak_memory_stats()

做资源核算时，用"手算的字节数"去对照 max_memory_allocated() 的实测值，是最好的自查方式——两者对不上，说明你漏算了某一项（通常是激活值或临时缓冲区）。

6. einops：给维度起名字

问题：位置维度太容易搞错

看这段传统 PyTorch 代码：

x = torch.ones(2, 2, 3)      # batch seq hidden
y = torch.ones(2, 2, 3)      # batch seq hidden
z = x @ y.transpose(-2, -1)  # batch seq seq

问题在哪？-2 和 -1 到底是什么？你必须在脑子里维护一张"第几个维度是什么"的表，而这张表只存在于注释里（还经常和代码不同步）。当张量变成 rank 4、rank 5，当代码里有十几处 transpose、permute、view，出错几乎是必然的。更糟的是，维度搞错往往不会报错——形状恰好匹配，只是算出来的东西没有意义。

einops 的解法是：用名字而不是位置来指代维度。它受爱因斯坦求和约定（Einstein summation notation, 1916）启发。三个核心函数：einsum、reduce、rearrange。（官方教程）

einsum：带记账功能的广义矩阵乘法

from einops import einsum, reduce, rearrange

x = torch.ones(3, 4)  # seq1 hidden
y = torch.ones(4, 3)  # hidden seq2

# 旧写法
z = x @ y

# einops 写法
z = einsum(x, y, "seq1 hidden, hidden seq2 -> seq1 seq2")

规则很简单：输入里出现、输出里没出现的维度，会被求和掉（上式中 hidden 被求和）。这正是矩阵乘法的定义 $z_{ij}=\sum_h x_{ih}y_{hj}$。

复杂一点的例子——带 batch 的注意力打分：

x = torch.ones(2, 3, 4)  # batch seq1 hidden
y = torch.ones(2, 3, 4)  # batch seq2 hidden

# 旧写法：必须记住 -2 -1 是什么
z = x @ y.transpose(-2, -1)   # batch seq1 seq2

# einops 写法：意图一目了然
z = einsum(x, y, "batch seq1 hidden, batch seq2 hidden -> batch seq1 seq2")

两边都有 batch 且输出保留了 batch，所以 batch 维不求和、按位对应（batched matmul）。hidden 在输出里消失，所以对它求和。你不需要写任何 transpose。

还可以用 ... 表示"任意多个前置维度"，让函数对任意 rank 都成立：

z = einsum(x, y, "... seq1 hidden, ... seq2 hidden -> ... seq1 seq2")

这在写通用模块时极其有用：同一段代码既能处理 [seq, hidden]，也能处理 [batch, head, seq, hidden]。

reduce：把某些维度归约掉

x = torch.ones(2, 3, 4)  # batch seq hidden

# 旧写法
y = x.sum(dim=-1)

# einops 写法
y = reduce(x, "... hidden -> ...", "sum")

支持 "sum"、"mean"、"max"、"min" 等。同样地，输出里没出现的维度就是被归约的维度——和 einsum 的心智模型完全一致。

rearrange：重排、拆分、合并维度

最典型的场景：某个维度实际上"打包"了两个维度（多头注意力里的 hidden = heads × head_dim），而你想只对其中一个做变换。

x = torch.ones(3, 8)   # seq total_hidden，其中 total_hidden = heads * hidden1
w = torch.ones(4, 4)   # hidden1 hidden2

# 1) 把 total_hidden 拆成 heads 和 hidden1
x = rearrange(x, "... (heads hidden1) -> ... heads hidden1", heads=2)
# 现在 x 的形状是 [3, 2, 4]

# 2) 只对 hidden1 做线性变换（heads 维自动作为"批"）
x = einsum(x, w, "... hidden1, hidden1 hidden2 -> ... hidden2")
# 形状 [3, 2, 4]

# 3) 把 heads 和 hidden2 合并回去
x = rearrange(x, "... heads hidden2 -> ... (heads hidden2)")
# 形状 [3, 8]

括号语法 (heads hidden1) 表示这个维度是两个维度按行优先展平得到的。拆分时必须给出其中一个的大小（heads=2），另一个由总长度推出。

7. FLOPs：算力的记账单位

讲完了"东西存在哪"，现在算"算它要多少代价"。

一次浮点运算（floating-point operation, FLOP）指一次基本运算，比如一次加法 x + y 或一次乘法 x * y。

先建立量级直觉

把最后一行和前两行比一比：8 张 H100 跑两周，只有 GPT-3 训练量的 3%，只有 GPT-4 的 0.05%。这就是"学术实验室 vs 前沿实验室"的算力鸿沟的具体数字。也正因如此，这门课的作业才要在小规模上做，靠 scaling law 外推。

矩阵乘法的 FLOPs：2mnp

这是全课最基础的一条公式，务必彻底搞懂它的来历。

课上的实测代码是这样的：

B = 16384  # Number of points（数据点数）
D = 32768  # Dimension of each point（输入维度）
K = 8192   # Number of outputs（输出维度）

x = torch.ones(B, D, device="cuda")
w = torch.randn(D, K, device="cuda")
y = x @ w

actual_num_flops = 2 * B * D * K       # = 8.796e12 FLOPs

$2 \times 16384 \times 32768 \times 8192 \approx 8.8\times 10^{12}$ FLOPs——一次矩阵乘法就快 9 TFLOPs。

实测：从时间反推 FLOP/s

import timeit

def benchmark(func, num_trials: int = 5) -> float:
    """返回执行 func 一次所需的秒数。"""
    if torch.cuda.is_available():
        torch.cuda.synchronize()      # 等待之前的 CUDA 任务结束

    def run():
        func()
        if torch.cuda.is_available():
            torch.cuda.synchronize()  # 等待本次任务真正跑完

    total_time = timeit.timeit(run, number=num_trials)
    return total_time / num_trials

actual_time = benchmark(lambda: x @ w)
actual_flop_per_sec = actual_num_flops / actual_time

硬件峰值算力：一定要看 dtype

每块 GPU 的规格书都给出峰值性能，而峰值严重依赖数据类型。课上代码里的查表函数是这样的：

def get_promised_flop_per_sec(dtype: torch.dtype) -> float:
    """返回该设备在给定 dtype 下的峰值 FLOP/s。"""
    properties = torch.cuda.get_device_properties("cuda")

    if "A100" in properties.name:
        if dtype == torch.float32:                     return 19.5e12
        if dtype in (torch.bfloat16, torch.float16):   return 312e12

    if "H100" in properties.name:
        if dtype == torch.float32:                     return 67.5e12
        if dtype in (torch.bfloat16, torch.float16):   return 1979e12 / 2
        # 1979 是带稀疏的数字，稠密只有一半

    if "B200" in properties.name:
        if dtype == torch.float32:                     return 75e12
        if dtype in (torch.bfloat16, torch.float16):   return 4.5e15 / 2

MFU：模型算力利用率

MFU（Model FLOPs Utilization，模型 FLOPs 利用率）的定义非常朴素：

（这个定义忽略了通信开销和其他 overhead——它们的影响会自动体现在"实测"那一项里。）

promised_flop_per_sec = get_promised_flop_per_sec(x.dtype)
mfu = actual_flop_per_sec / promised_flop_per_sec

经验判断：MFU $\geq$ 0.5 就已经相当好了。 前沿实验室的大规模预训练通常报告 0.4–0.55；能上 0.6 说明工程做得很扎实。

实测方法就是上面那套：跑一步训练，测时间，用 $6ND$ 算出"模型 FLOPs"，除以时间得到实测 FLOP/s，再除以峰值。

那么问题来了：为什么 MFU 不接近 1？ 明明矩阵乘法的 FLOPs 是确定的，硬件峰值也是确定的。要回答这个问题，必须看清 GPU 上计算到底是怎么发生的——这就是下一节。

8. 算术强度与 roofline：为什么打不满算力

做一件计算，物理上要走三步：

1. 把输入从内存搬到加速器（memory → accelerator）
2. 执行计算
3. 把输出从加速器搬回内存

所以耗时由两个参数共同决定：

• 加速器速度：$\text{h100\_flop\_per\_sec} = 1979\times10^{12}/2 = 9.895\times 10^{14}$ FLOP/s
• 内存带宽：$\text{h100\_bytes\_per\_sec} = 3.35\times 10^{12}$ B/s

假设通信和计算能完美重叠（理想情况），那么

两个"强度"

比时间更好用的是一个无量纲的比值。定义：

295 这个数字要记住：H100 在 bf16 下，每从显存读一个字节，就"有资格"做 295 次浮点运算。你的算子做不到这么多，就是在浪费芯片。

下面按算术强度从低到高，把五种典型算子过一遍。所有例子都用 bf16（2 字节/元素）。

例 1：ReLU —— 强度 0.25

n = 1024 * 1024
x = torch.ones(n, dtype=torch.bfloat16, device="cuda")
y = torch.relu(x)

bytes = (2 * n) + (2 * n)   # 读 x，写 y（bf16 每个元素 2 字节）
flops = n                   # n 次比较

communication_time = bytes / h100_bytes_per_sec   # ≈ 1.25e-6 s
computation_time   = flops / h100_flop_per_sec    # ≈ 1.06e-9 s

arithmetic_intensity = flops / bytes              # = 0.25

算术强度 0.25，加速器强度 295。差了一千多倍！通信时间比计算时间长约 1180 倍——GPU 有 99.9% 的时间在等数据。ReLU 是彻头彻尾的 memory-bound。

例 2：GELU —— 强度 5

y = F.gelu(x)
# GELU(x) = 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))

bytes = (2 * n) + (2 * n)   # 还是读 x、写 y
flops = 20 * n              # tanh 可用多项式近似，估 20 次运算

arithmetic_intensity = flops / bytes    # = 5

GELU 每搬一个字节干的活比 ReLU 多得多（20 倍），强度提到 5——但离 295 还差 60 倍，依然是 memory-bound。

例 3：点积（向量 · 向量）—— 强度 0.5

x = torch.ones(n, dtype=torch.bfloat16, device="cuda")
w = torch.ones(n, dtype=torch.bfloat16, device="cuda")
y = x @ w

bytes = (2 * n) + (2 * n) + 2   # 读 x、读 w、写一个标量 y
flops = 2 * n - 1               # n 次乘法，n-1 次加法

arithmetic_intensity = flops / bytes   # ≈ 0.5

Memory-bound。

例 4：矩阵-向量乘法 —— 强度 ≈ 1

n = 1024
x = torch.ones(n, dtype=torch.bfloat16, device="cuda")       # 向量
w = torch.ones(n, n, dtype=torch.bfloat16, device="cuda")    # 矩阵
y = x @ w

bytes = (2 * n) + (2 * n * n) + (2 * n)   # 读 x、读 w、写 y
flops = n * (2 * n - 1)                   # n 个点积

arithmetic_intensity = flops / bytes      # ≈ 1

算一下：$\dfrac{2n^2}{2n^2 + 4n}\approx 1$。还是 memory-bound！ 直觉是：矩阵 $W$ 有 $n^2$ 个元素要读进来，但每个元素只被用了一次。搬运的成本完全没被摊薄。

例 5：矩阵-矩阵乘法 —— 强度 n/3，终于 compute-bound

n = 1024
x = torch.ones(n, n, dtype=torch.bfloat16, device="cuda")
w = torch.ones(n, n, dtype=torch.bfloat16, device="cuda")
y = x @ w

bytes = (2 * n * n) + (2 * n * n) + (2 * n * n)   # 读 x、读 w、写 y
flops = n * n * (2 * n - 1)                      # n^2 个点积

arithmetic_intensity = flops / bytes   # ≈ 2n^3 / 6n^2 = n / 3

关键在于：搬运量是 $O(n^2)$，而计算量是 $O(n^3)$。矩阵越大，每个搬进来的元素被复用的次数越多，强度线性增长。

$n=1024$ 时 $I_{\text{arith}} = 341 > 295$，终于 compute-bound 了！ 而且刚刚好越线——反解一下：

Percy 的总结：

• 只要矩阵够大，我们就是 compute-bound（能吃满加速器）。
• 训练 Transformer 的主体就是大矩阵乘法——所以训练能有不错的 MFU。
• 矩阵-向量乘法是推理时发生的事，所以推理是 memory-bound。
• 注意：算术强度和加速器强度都依赖精度。换到 fp32，$I_{\text{acc}} = 67.5\times10^{12}/3.35\times10^{12}\approx 20$，门槛低了很多（但绝对算力也低了 15 倍）；换到 fp8，$I_{\text{acc}}$ 翻倍到约 590（且每元素只有 1 字节，$I_{\text{arith}}$ 也会变），门槛更高。

Roofline 图

把上面所有分析画成一张图，就是经典的 roofline plot（JAX Scaling Book 的 roofline 章节有很好的可视化）：

• 横轴：算术强度（FLOP/byte）。每个具体计算是横轴上的一个点。
• 纵轴：达到的性能（FLOP/s）。
• 每条分段线性的折线代表一种硬件：左边是斜率为"内存带宽"的上升段（memory-bound 区，性能 = 强度 × 带宽），右边是水平段（compute-bound 区，性能 = 峰值算力）。
• 拐点（kink）就是加速器强度——从 memory-bound 过渡到 compute-bound 的分界，H100 bf16 下约 295 FLOP/byte。

这张图把 MFU 也解释清楚了：

也就是说，在理想的重叠假设下，MFU 上不去的根本原因就是算术强度不够。ReLU 的 MFU 上界是 $0.25/295\approx 0.00085$；$1024\times1024$ 矩阵乘法的上界是 1。真实训练混合了两类算子，加上通信、kernel launch、非重叠部分，最终落在 0.4–0.5 附近。

9. 反向传播的 FLOPs：6ND 的完整推导

到目前为止我们只做了前向。现在把反向加进来，得到本课最常用的一个公式。

先热身：autograd 在做什么

用一个极简的线性模型 $y = \tfrac{1}{2}(x\cdot w - 5)^2$：

# 前向：计算 loss
x = torch.tensor([1., 2, 3])
w = torch.tensor([1., 1, 1], requires_grad=True)   # 想要它的梯度
pred_y = x @ w                     # = 6
loss = 0.5 * (pred_y - 5).pow(2)   # = 0.5

# 反向：计算梯度
loss.backward()
assert torch.equal(w.grad, torch.tensor([1., 2, 3]))

手推验证：$\dfrac{\partial \text{loss}}{\partial \hat y} = \hat y - 5 = 1$，再由 $\hat y = x\cdot w$ 得 $\dfrac{\partial \text{loss}}{\partial w} = 1 \cdot x = [1,2,3]$。✓

要点：requires_grad=True 让 PyTorch 在前向时记录计算图并保留中间结果——这就是激活值显存的来源。backward() 沿图反向走一遍链式法则。

数清楚反向的 FLOPs

用一个两层线性网络来数：

B = 1024  # Number of points（数据点数）
D = 256   # Dimension（维度）

x  = torch.ones(B, D, device="cuda")
w1 = torch.randn(D, D, device="cuda", requires_grad=True)
w2 = torch.randn(D, D, device="cuda", requires_grad=True)

# 前向
h1 = einsum(x,  w1, "batch in, in out -> batch out")   # 等价于 x @ w1
h2 = einsum(h1, w2, "batch in, in out -> batch out")   # 等价于 h1 @ w2
loss = (h2.mean() - 0)**2      # 随便回归到 0

# 反向
h1.retain_grad()   # 只是为了方便调试时查看
h2.retain_grad()
loss.backward()

聚焦第二层：$h_2 = h_1 W_2$，其中 $h_1\in\R^{B\times D}$，$W_2 \in \R^{D\times D}$。

前向的 FLOPs（用上一节的 $2mnp$）：

反向要算两样东西：

课上用 einsum 把这两个式子直接写了出来并和 autograd 对拍——非常值得自己跑一遍，这是理解反向传播最快的方式：

h1_grad = einsum(h2.grad, w2, "batch out, in out -> batch in")
assert torch.allclose(h1.grad, h1_grad)

w2_grad = einsum(h2.grad, h1, "batch out, batch in -> in out")
assert torch.allclose(w2.grad, w2_grad)

num_backward_flops = (2 * B * D * D) + (2 * B * D * D)   # = 4 * B * D * D

注意 einsum 的写法把这两个转置藏进了名字里："batch out, in out -> batch in" 里 out 被求和、in 被保留，自然就是 $W_2^\top$ 的效果，你根本不用去想哪个维度要 transpose。

推广到整个网络：6ND

上面只是 $W_2$ 一个参数矩阵。网络里每个参数矩阵都要走同一套流程，把它们加起来：

为什么前向是 $2ND$？回到 $2mnp$：一个 $[D_{\text{in}}, D_{\text{out}}]$ 的权重矩阵有 $n_W = D_{\text{in}}D_{\text{out}}$ 个参数，作用在 $B$ 个数据点上要 $2\,B\,D_{\text{in}}D_{\text{out}} = 2\,B\,n_W$ FLOPs。对所有权重求和，就是 $2\,B\,N$。 这个"每参数每 token 2 FLOPs"的形式之所以成立，正是因为参数量本身就等于矩阵元素个数。

10. 模型、初始化、优化器：把显存逐项列出来

一个最小的深度网络

import math
import torch.nn.functional as F
from torch import nn

class Block(nn.Module):
    """一个线性变换 + ReLU 非线性。"""
    def __init__(self, dim: int):
        super().__init__()
        # 关键：除以 sqrt(dim)
        self.weight = nn.Parameter(torch.randn(dim, dim) / math.sqrt(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x @ self.weight   # 线性
        x = F.relu(x)         # 激活
        return x


class DeepNetwork(nn.Module):
    """把 dim 维向量映射到 dim 维向量。"""
    def __init__(self, dim: int, num_layers: int):
        super().__init__()
        self.layers = nn.ModuleList([Block(dim) for _ in range(num_layers)])

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        for layer in self.layers:
            x = layer(x)
        return x


D, L = 8, 3
model = DeepNetwork(dim=D, num_layers=L).to(device)

def get_num_parameters(model: nn.Module) -> int:
    return sum(p.numel() for p in model.parameters())

assert get_num_parameters(model) == (D * D) * L

两个 PyTorch 机制值得注意：

• nn.Parameter 把一个张量登记成"模型参数"：它会自动 requires_grad=True，会出现在 model.parameters() 和 state_dict() 里，.to(device) 时会跟着搬。
• nn.ModuleList 而不是普通 Python list：普通 list 里的子模块不会被注册，参数会"消失"（不进 parameters()、不上 GPU、不被优化）。这是新手最容易踩的坑之一。

为什么初始化要除以根号 d

torch.randn(dim, dim) / math.sqrt(dim) 里的 $1/\sqrt{d}$ 不是随手写的。

实践中用截断正态（truncated normal）

真实实现里通常不是直接用 randn，而是用截断正态分布：

std = 1.0 / math.sqrt(d)          # 或 sqrt(2 / (d_in + d_out))
w = torch.empty(d_in, d_out)
nn.init.trunc_normal_(w, mean=0.0, std=std, a=-3 * std, b=3 * std)

理由：

• 掐掉极端离群值。 一个 $d^2 = 10^6$ 个元素的权重矩阵里，正态分布必然出现若干个 $4\sigma$、$5\sigma$ 的值。这些异常大的权重在训练最初几十步会制造异常大的激活和梯度，是早期 loss spike 的常见诱因。截断在 $\pm 3\sigma$ 就把这个尾巴剪掉了。
• 数值可控。 有明确上下界，配合低精度（bf16/fp8）时更安全。

常见的 std 选择有两类：1/sqrt(d_in)（只保前向方差，即上面的推导）和 sqrt(2/(d_in + d_out))（Xavier，前向反向折中）。对深层残差网络还会额外把某些层再除以 $\sqrt{2L}$，以免残差流的方差随深度累积。

优化器：从 SGD 到 AdamW 的家谱

Percy 用一条清晰的链条把常用优化器串起来：

AdaGrad 从"$\sum g^2$"到 RMSProp 的"指数平均"这一步很关键：AdaGrad 的分母只增不减，学习率会单调衰减到 0，对长时间训练不友好；指数平均让它能"忘记"很久以前的梯度。

课上手写了一个 AdaGrad（作业 1 会让你按同样的模板写 AdamW）：

class AdaGrad(torch.optim.Optimizer):
    def __init__(self, params: Iterable[nn.Parameter], lr: float = 0.01):
        super().__init__(params, dict(lr=lr))

    def step(self):
        for group in self.param_groups:
            lr = group["lr"]
            for p in group["params"]:
                state = self.state[p]      # 每个参数对应一份状态字典
                grad = p.grad.data

                # 取出累积的平方梯度 g2 = sum_{i<t} g_i^2
                g2 = state.get("g2", torch.zeros_like(grad))

                # 更新优化器状态
                g2 += torch.square(grad)
                state["g2"] = g2

                # 更新参数
                p.data -= lr * grad / torch.sqrt(g2 + 1e-5)

模板要点：self.param_groups 是参数分组（可以给不同组不同学习率），self.state[p] 是每个参数张量一份的状态字典——优化器状态的显存就是从这里来的，它和参数张量同形状。

显存的完整清单

现在把所有项摆上桌。以 $D=4$、$L=3$、$B=2$ 的玩具网络为例（参数量 $N = D^2 L = 48$）：

num_parameters = D * D * L                    # 48

parameter_memory       = 2 * num_parameters   # bf16，2 字节  → 96 B
gradient_memory        = 2 * num_parameters   # bf16，2 字节  → 96 B
optimizer_state_memory = 4 * num_parameters   # fp32，AdaGrad 一份状态 → 192 B
activation_memory      = 2 * (B * D * L)      # bf16，每层每样本一份 → 48 B

total_memory = (parameter_memory + gradient_memory
                + optimizer_state_memory + activation_memory)   # 432 B

为什么优化器状态用 fp32？课上给了明确理由：它要在成千上万步里累积平均（"accumulating averages over powers over many steps"），而 bf16 只有 7 位尾数——当动量已经积累到 $1.0$ 时，一个 $0.001$ 的新增量在 bf16 里会被直接丢弃（$1.0 + 0.001 = 1.0$）。参数和梯度是"一次性"的量，精度低一点无所谓；优化器状态是"长期记忆"，必须保精度。

回到开头的问题：8 张 H100 能训多大？

再看几个变体（同样忽略激活）：

训练比推理贵 6 倍显存——而且这还没算激活。这张表也解释了为什么会有 8-bit Adam、Adafactor 这类"省优化器状态"的工作。

一步训练的算力

flops = 6 * B * num_parameters     # 6 * 2 * 48 = 576 FLOPs

对 Transformer 来说核算更复杂，但思路完全一样：把每个矩阵乘法的三个维度写出来，套 $2mnp$，再乘 3（前向 + 反向）。作业 1 就要求你对 Transformer 做完整的 FLOPs 和显存核算。

两篇很好的参考：Transformer 训练的显存分析、Transformer 的 FLOPs 计算。

11. 一个完整的极简训练循环

把前面所有零件拼起来。一个训练脚本永远是这五块：数据 → 模型 → 前向 → 反向 → 更新，外加 checkpoint。

骨架

# 真实的线性函数，权重是 (0, 1, 2, ..., D-1)
D = 16
true_w = torch.arange(D, dtype=torch.float32, device=device)

# 数据加载：生成 (x, y) 对
B = 4
def get_batch() -> tuple[torch.Tensor, torch.Tensor]:
    x = torch.randn(B, D).to(device)
    true_y = x @ true_w
    return (x, true_y)

# 模型和优化器
L = 2
model = DeepNetwork(dim=D, num_layers=L).to(device)
optimizer = AdaGrad(model.parameters(), lr=0.01)

# 训练！
num_train_steps = 3
for t in range(num_train_steps):
    # 取数据
    x, y = get_batch()

    # 前向（算 loss）
    pred_y = model(x).mean()
    loss = F.mse_loss(pred_y, y)

    # 反向（算梯度）
    loss.backward()

    # 更新参数
    optimizer.step()
    optimizer.zero_grad(set_to_none=True)

真实的数据加载：np.memmap

玩具例子里数据是现场生成的。真实预训练的数据是一个几 TB 的 token 数组，不可能读进内存。

先算账：分词后每个 token 是一个整数。词表通常小于 65536，所以用 uint16 存，2 字节/token。

解决办法是 np.memmap：把磁盘上的文件映射到虚拟地址空间，像操作数组一样索引它，由操作系统按需分页调入。内存占用只和你实际访问的部分有关。

import numpy as np

# 训练前把 token 一次性写成一个扁平的 uint16 二进制文件
# tokens.astype(np.uint16).tofile("train.bin")

data = np.memmap("train.bin", dtype=np.uint16, mode="r")   # 不读进内存

def get_batch(data: np.memmap, batch_size: int, seq_len: int,
              device: str) -> tuple[torch.Tensor, torch.Tensor]:
    # 在整个语料里随机采 batch_size 个起点
    starts = np.random.randint(0, len(data) - seq_len - 1, size=batch_size)

    # 输入是 [i, i+seq_len)，标签是右移一位的 [i+1, i+1+seq_len)
    x = np.stack([data[i     : i + seq_len    ] for i in starts])
    y = np.stack([data[i + 1 : i + 1 + seq_len] for i in starts])

    # uint16 -> int64（embedding 查表需要 long）
    x = torch.from_numpy(x.astype(np.int64))
    y = torch.from_numpy(y.astype(np.int64))

    if device.startswith("cuda"):
        # pin_memory + non_blocking：让 H2D 拷贝和计算重叠
        x = x.pin_memory().to(device, non_blocking=True)
        y = y.pin_memory().to(device, non_blocking=True)
    else:
        x, y = x.to(device), y.to(device)
    return x, y

几个要点：

• 随机起点采样而不是顺序遍历：省掉了维护 epoch 和 shuffle 的复杂度，在超大语料上两者统计上等价。
• 标签就是右移一位的输入——这是自回归语言建模的全部秘密，不需要单独存标签。
• 页缓存（page cache）：第一次访问某段数据会触发磁盘 I/O，之后 OS 会缓存。所以真实训练要保证磁盘吞吐跟得上。
• pin_memory() + non_blocking=True：锁页内存才能走 DMA 异步拷贝，让下一个 batch 的传输和当前 batch 的计算重叠，否则 GPU 会周期性空等。

Checkpoint

一次训练要跑几个月，机器一定会挂。checkpoint 必须同时存模型、优化器状态、当前步数——只存模型是不够的，Adam 的动量丢了会导致恢复后 loss 抖动。

def save_checkpoint(model: nn.Module, optimizer: torch.optim.Optimizer,
                    iteration: int, out: str):
    torch.save({
        "model": model.state_dict(),
        "optimizer": optimizer.state_dict(),
        "iteration": iteration,
    }, out)

def load_checkpoint(src: str, model: nn.Module,
                    optimizer: torch.optim.Optimizer) -> int:
    ckpt = torch.load(src, map_location="cpu")
    model.load_state_dict(ckpt["model"])
    optimizer.load_state_dict(ckpt["optimizer"])
    return ckpt["iteration"]

12. 用计算换显存：梯度累积与激活重计算

前面把显存拆成了两类：模型侧（参数/梯度/优化器状态，$12N$ 字节） 和 批次侧（激活值）。这一节讲怎么压缩第二类。

梯度累积（gradient accumulation）

动机：大 batch 有利于训练稳定性（梯度噪声小、更容易用大学习率、更好地打满 GPU）。但激活显存随 batch 线性增长，很容易 OOM。

B = 64     # Batch size
D = 1024   # Dimensionality
L = 16     # Number of layers

activation_memory = 2 * B * D * L   # bf16 → 2,097,152 字节 ≈ 2 MB

梯度累积的想法非常简单：

• 把大 batch 切成若干个 micro batch，逐个做前向 + 反向；
• 累加梯度，不清零（正好利用了 PyTorch 梯度默认累加的行为）；
• 每累积 batch_size / micro_batch_size 次，才调用一次 optimizer.step() 并清零。

micro_batch_size = B / 4                            # 16
activation_memory = 2 * micro_batch_size * D * L    # → 524,288 字节 ≈ 0.5 MB

激活显存降到 1/4，而数学上等价于用 batch 64 训练（梯度是各 micro batch 梯度之和）。

accum_steps = 4
for t in range(num_train_steps):
    for micro in range(accum_steps):
        x, y = get_batch(micro_batch_size)
        loss = compute_loss(model, x, y) / accum_steps   # 注意除以 accum_steps
        loss.backward()          # 梯度自动累加，不清零
    optimizer.step()
    optimizer.zero_grad(set_to_none=True)

激活重计算（activation checkpointing）

先看清问题的本质：

• 训练：反向传播需要每一层的激活（回忆 $\partial\mathcal{L}/\partial W_2 = h_1^\top\,\partial\mathcal{L}/\partial h_2$，必须有 $h_1$），所以所有层的激活都要存着。
• 推理：不算梯度，只需要当前层的激活，用完就扔，显存 $O(1)$。

B, D, L = 64, 1024, 16
activation_memory = 2 * B * D * L   # ≈ 2 MB（bf16）

x = torch.randn(B, D, device=device, requires_grad=True)
model = DeepNetwork(dim=D, num_layers=L).to(device)
memory = get_max_memory_usage(lambda: model(x).sum().backward())

能不能减少？可以。激活重计算 = 梯度检查点（gradient checkpointing）= 重物化（rematerialization），三个名字一回事：

课上的对比示意（$g_i$ 表示层内部的中间量，$h_i$ 表示层输出）：

存所有激活:  x g1 h1 g2 h2 g3 h3 g4 h4
激活重计算:  x    h1    h2    h3    h4

PyTorch 里只要把 layer(x) 换成 torch.utils.checkpoint.checkpoint(layer, x)：

class DeepNetworkCheckpointed(nn.Module):
    """和 DeepNetwork 一样，但开了激活重计算。"""
    def __init__(self, dim: int, num_layers: int):
        super().__init__()
        self.layers = nn.ModuleList([Block(dim) for _ in range(num_layers)])

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        for layer in self.layers:
            # 关键：只在检查点处保存激活，其余的反向时重算
            x = torch.utils.checkpoint.checkpoint(layer, x)
        return x

多久打一个检查点？

这是一个漂亮的时间-空间权衡题。设网络有 $L$ 层：

实践中的常见做法：以一个 Transformer block 为粒度打检查点（即 $k=1$，$O(L)$ 个检查点但每个 block 内部的中间激活全部重算）。这已经能省掉大部分激活显存，而额外计算约为一次前向，即总计算从 $6ND$ 变成 $8ND$——多约 33%。

本讲小结

Percy 在课上的六条总结：

1. 一切都是张量上的操作——参数、梯度、激活值、优化器状态、数据。
2. einops 是思考张量操作更好的方式：给维度起名字，而不是数 -1 -2。
3. 每步训练 $6 \times (\text{数据点数}) \times (\text{参数量})$ FLOPs。
4. 算术强度 / roofline 分析：判断一个计算是 compute-bound 还是 memory-bound。
5. 矩阵乘法是 compute-bound，逐元素操作是 memory-bound。
6. 梯度累积、激活重计算：降低显存以支持更大的 batch。

速查表

可以马上拿去用的三个习惯

1. 看到任何模型配置，先算 $6ND$ 和 $12N$。 前者告诉你要多少卡·天，后者告诉你要多少卡。
2. 看到任何算子，先算算术强度。 低于 295（H100/bf16）就别指望它快，优化方向是减少字节而不是减少 FLOPs。
3. 看到任何 OOM，先把显存拆成五项（参数/梯度/优化器/激活/临时缓冲），逐项对照 torch.cuda.max_memory_allocated()。手算和实测对不上的那一项，就是你的 bug。

下一讲会把这套核算方法用到具体的 Transformer 架构上：每个架构选择（Pre-LN vs Post-LN、RMSNorm、SwiGLU、RoPE、GQA）都会同时改变 FLOPs、显存和稳定性，而你现在已经有工具去评估它们了。

附录：延伸阅读

数值精度

• Mixed Precision Training (2017) — 混合精度训练的奠基论文，loss scaling、fp32 主权重这些标准做法都出自这里。
• FP8 Formats for Deep Learning (2022) — E4M3 / E5M2 两种 fp8 格式的标准化提案，解释了为什么前向和反向要用不同格式。
• Introducing NVFP4 (2025) — 4 比特浮点 + 分块缩放因子的设计，看"极限压缩"能走到哪里。
• Nemotron 3 Super Technical Report (2026) — 用 NVFP4 完成的真实大规模训练，是低精度训练可行性的最新证据。
• PyTorch AMP 文档 — 哪些算子会被自动降精度、哪些不会，值得完整读一遍白名单。

张量与 einops

• Einops 教程 — 半小时读完，之后写模型的心智负担会显著下降。
• DeepSeek v3.2 权重清单 — 用真实模型的张量列表去练"看形状猜结构"。
• DeepSeek v3.2 (2025) — 对应的技术报告。

算力与 roofline

• How to Scale Your Model：Roofline 章节 — 本讲算术强度部分最好的配套材料，有大量可交互的图。
• Transformer FLOPs 详解 — 把 $6ND$ 的每一项按 Transformer 的实际结构拆开，含注意力修正项。
• Transformer 训练显存详解 — 参数/梯度/优化器/激活逐项核算，和作业 1 的要求高度对应。
• Scaling Laws for Neural Language Models (2020) — 附录里有 $6ND$ 及注意力修正项的经典推导。

优化器

• Adaptive Subgradient Methods (AdaGrad, 2011) — 本讲手写实现的那个优化器的原论文。
• Adam (2014) — RMSProp + momentum，两个动量状态的来源。
• Decoupled Weight Decay Regularization (AdamW, 2017) — 今天训练 LLM 的默认优化器，解释了为什么 Adam 里的 L2 正则是错的。

省显存的技术

• GPipe (2018) — 流水线并行，同时也是激活重计算在大规模训练中的早期系统性应用。
• Reducing Activation Recomputation in Large Transformer Models (2022) — 选择性重计算：只重算"存着贵、算着便宜"的那部分。
• ZeRO (2019) — 把参数/梯度/优化器状态分片到多卡，直接攻击那 $12N$ 字节。
• FlashAttention (2022) — 算术强度思想的教科书级应用：不物化 $L\times L$ 注意力矩阵，把 memory-bound 的注意力变回 compute-bound。

整体图景

• Attention Is All You Need (2017) — 所有核算的对象。
• Training Compute-Optimal Large Language Models (Chinchilla, 2022) — 在 $6ND$ 的预算约束下，$N$ 和 $D$ 该怎么分配。本讲的公式是它的前提。