9.1 问题设定：超过单卡容量

第 1 讲到第 8 讲训练的都是「一张卡装得下」的模型：前向、反向、更新，三步循环。这一节要先把问题量化清楚——因为「模型太大了」这句话里藏着至少三种完全不同的困难，而它们的处方互不相同。

具体地说，这一节回答三个问题：

显存里到底放了什么？ 参数量只占其中一小部分，真正的大头是优化器状态与激活——算清这笔账，才知道该省哪一项。
除了显存，还有哪几堵墙？ 显存、算力、通信三者互相牵制，只看显存会得出错误的方案。
「放不下」有几种情形？ 模型状态放不下、激活放不下、单个层放不下——分别对应不同的切分手段。

9.1.1 一次训练步骤要在显存里放什么

训练不是「把模型放进去」就完了。一次完整的训练步骤里，以下五类张量必须同时驻留：

编号	项	形状	为什么必须在场
①	参数（weights）	与 $N$ 同形	前向要用；反向还要用它算梯度
②	梯度（gradients）	与 $N$ 同形	反向产生、更新消耗；在所有 micro-batch 累加期间一直在场
③	优化器状态（optimizer state）	与 $N$ 同形 × 2	Adam 族的一阶矩 $m$ 、二阶矩 $v$ ，每一步更新都要读改写
④	主权重（master weights）	与 $N$ 同形	低精度参数上做小步更新会被舍入吃掉，故另存一份高精度副本
⑤	激活（activations）	与批大小 × 序列长度 × 层数成正比	反向传播需要前向的中间结果（1.5）

前四项称为模型状态（model states）：它们的规模只与参数量 $N$ 有关，与批大小无关。第五项是激活：它不随参数量直接变化，却随「一次喂进去多少 token」线性增长。

定义 9.1.1 · 模型状态与激活

对含 $N$ 个可训练参数的模型，一次训练步骤的显存占用可划分为两部分：

模型状态 $M_{\text{state}}$ ：参数、梯度、优化器状态、主权重。四者都与 $N$ 成正比，只要训练在继续就必须常驻；
激活 $M_{\text{act}}$ ：前向产生的中间张量。它与「单个 micro-batch 的 token 数」 $B_{\text{micro}}\times T$ 、层数 $L$ 成正比，在反向用完某一层的激活后即可释放。

总占用（不含通信缓冲与碎片）为 $M \approx M_{\text{state}} + M_{\text{act}} + M_{\text{comm}} .$

其中 $M_{\text{comm}}$ 是集合通信的接收缓冲、梯度分桶等临时空间，通常由框架管理，量级介于前两者之间。

为什么「账要分成两半算」

把显存想成一间房：模型状态是搬不走的家具（从第一步到最后一步都在场），激活是一次宴会的客人（来一波走一波，但人多了照样挤爆房间）。

这个区分决定了两类处方的分工：

房间小、家具大 → 换更薄的家具（低精度、8-bit 优化器）或把家具分散到多间房（ZeRO / FSDP 分片、张量并行、流水线并行）；
家具不大但客人太多 → 换一批少请几个（减小 micro-batch、梯度累积）或让客人快点走（激活重计算）。

搞混这两个方向的典型后果是：为了省激活去开 ZeRO-3（它分的是模型状态），或者为了省模型状态去减小批大小（批大小根本不影响模型状态）。

9.1.2 混合精度 Adam 下的 $16N$ 规则

现在把账算到字节级。约定：混合精度训练（参数与梯度用 16 位，优化器相关量用 fp32），优化器为 Adam / AdamW。

五账单的逐项推导（混合精度 Adam）

设参数量为 $N$ ，把每一项写成「每参数字节数 × $N$ 」：

① 参数：前向与反向要读参数，用 bf16（或 fp16）存储 → $\mathbf{2N}$ 字节。

② 梯度：反向产生，与参数同精度 → $\mathbf{2N}$ 字节。（若梯度用 fp32，这项变 $4N$ ；部分框架提供「bf16 梯度 + fp32 累加」的折中。）

③ 优化器状态：Adam 对每个参数维护一阶矩 $m$ 与二阶矩 $v$ 。它们必须是 fp32： $m$ 、 $v$ 的数值尺度随梯度方差的估计变化，用 16 位会让长期累积的动量被舍入噪声污染 → $4N+4N=\mathbf{8N}$ 字节。

这一项是账本里最贵的单项：仅 $m$ 、 $v$ 就占了总账的一半。换成 SGD（无动量）时它降为 0，这也是「SGD 省显存」这句话的全部含义。

④ 主权重：更新公式为 $w_{\text{fp32}} \leftarrow w_{\text{fp32}} - \eta\cdot\hat g$ ，因为 $\eta\hat g$ 的量级常小于 bf16 在该权重附近的间隔（bf16 每参数只有 8 位尾数），直接更新 16 位权重会出现「加了等于没加」 → $\mathbf{4N}$ 字节。

⑤ 激活：与批大小、序列长度、层数成正比，见 9.1.3；不计入这四账的常数项。

合计： $M_{\text{state}} \approx 2N+2N+8N+4N=\mathbf{16N}\ \text{字节} .$

推理期只需 ①（低精度参数）+ KV cache：同一模型在推理时是 $2N$ 字节量级，这是「训练比推理难得多」在显存上的定量来源。

16N

的正确读法：三项在参数上，两项在优化器上

记住两个近似比例：

优化器相关占四分之三： $8N$ （ $m,v$ ）+ $4N$ （主权重）= $12N$ ，占 $16N$ 的 $75\%$ 。而参数与梯度合计只有 $4N$ ，占四分之一。
fp32 相关占七成五：优化器状态与主权重都是 fp32， $12N$ 全部来自「为了数值稳定而保留高精度」。

由此可以预判所有省显存手段的排名：动优化器状态收益最大（8-bit 优化器、ZeRO-1 分片、换 SGD），动精度次之（混合精度），动参数最后（ZeRO-3 把参数也分片，省的是那 $2N$ 中的一部分）。

例 9.1.1 · 7B 模型的显存账

一个 $N=7\times10^{9}$ （70 亿参数）的语言模型，用混合精度 + AdamW 训练。求：

（1）模型状态需要多少显存？（2）一张 80 GB 的训练卡装得下吗？（3）改用 8-bit 优化器后是多少？（4）用 8 张卡做 ZeRO-3 全分片后每卡多少？

（按 $1\,\text{GB}=10^{9}$ 字节估算，便于心算；写成 $2^{30}$ 时结果约小 7%，量级结论不变。）

解

（1）模型状态。 由五账单的逐项推导， $M_{\text{state}}=16N=16\times7\times10^{9}=112\times10^{9}\ \text{字节}=112\ \text{GB}.$ 逐项拆开核对：参数 $2\times7=14$ GB，梯度 $14$ GB，主权重 $28$ GB， $m$ 与 $v$ 各 $28$ GB；合计 $14+14+28+28+28=112$ GB。✔

（2）80 GB 卡。 $112\ \text{GB} \gt 80\ \text{GB}$ ：只装模型状态就已经超了 40%，再加激活、（若用数据并行）通信缓冲与显存碎片，实际需要 130 GB 以上。结论：这张卡连「还没开始算」都过不去。

（3）8-bit 优化器。 把 $m$ 、 $v$ 压到 1 字节/参数（见 9.5.4），主权重仍留 fp32： $M_{\text{state}}=(2+2+1+1+4)N=10N=70\ \text{GB}.$ 比 112 GB 省 $37.5\%$ 。但仍高于 80 GB 减去激活后的可用空间——省显存与分片要配合使用，没有哪一招单独解决 7B。

（4）ZeRO-3 / FSDP，8 卡。 把四项模型状态全按 $k=8$ 分片： $M_{\text{state}}=16N/k=14$ GB/卡。于是 80 GB 的卡上还剩约 60 GB 给激活、通信缓冲与碎片——「放不下」在这一步变成「放得下」。代价是每算一层都要把该层参数临时聚合（all-gather），通信量上升，见 9.2.5 与 9.6.1。 $\blacksquare$

∎

9.1.3 激活账：OOM 最常见的真凶

模型状态可以分片、可以降精度，但激活在同一张卡上随「一次处理的 token 数」线性增长，而且分片手段（ZeRO-1/2/3）完全不作用于它。这就是为什么很多人把优化器分片做到极限后仍然 OOM。

对标准 Transformer，逐层的激活主要有这几类（ $B$ 为 micro-batch、 $T$ 为序列长度、 $h$ 为隐藏维、 $L$ 为层数）：

激活项	形状	是否需要保存到反向	备注
LayerNorm 输入 / 输出	$B\times T\times h$	需要	每层 2 份
注意力 $Q,K,V$	$3\times B\times T\times h$	需要	张量并行下按头切分
注意力概率矩阵	$B\times n_h\times T\times T$	取决于实现	显存的主要不确定项
注意力输出投影输入	$B\times T\times h$	需要	行并行时是部分和
MLP 隐层（如 $4h$ ）	$B\times T\times 4h$	需要	SwiGLU 有两个并列分支
Dropout 掩码	与作用对象同形	需要	可按 1 bit/token 存储

关键结论：若注意力矩阵被实体化，它的开销是 $O(B,n_h,T^{2})$ ，随序列长度平方增长——这条曲线是所有「长上下文训不动」问题的共同根源。FlashAttention 类实现（2.6）用分块与在线 softmax 避免了实体化这一项，代价是反向要重算注意力；激活重计算（9.5.3）则是把整个前向的激活都不存。

例 9.1.2 · 给一次 8 卡训练估激活账

配置： $L=32$ 层， $h=4096$ ， $n_h=32$ ， $B_{\text{micro}}=4$ ， $T=2048$ ，bf16（2 字节/元素），不开张量并行。按「每层每 token 的每个 hidden 元素约 34 字节」的经验口径（该口径来自 Megatron-LM 的激活量估算：把 LayerNorm 输入输出、QKV、注意力输出、MLP 隐层等中间张量按精度折算后求和，得到一个与 $h$ 无关的每元素系数）估算激活总量。

解

逐步算。

单个 micro-batch 的 token 数： $B_{\text{micro}}T=4\times2048=8192$ 。

按经验口径，激活总量 ≈ $34\ \text{字节} \times (\text{token 数}) \times h \times L$ ——注意 $h$ 这一项不能漏（34 是「每个 hidden 元素」的字节数，不是「每个 token」的字节数）： $M_{\text{act}} \approx 34\times8192\times4096\times32 \approx 3.7\times10^{10}\ \text{字节} \approx 34\ \text{GB}.$

对账（为什么 34 这个系数与 $h$ 无关）：LayerNorm 的输入与输出各 $BTh\times2$ 字节， $Q,K,V$ 共 $3BTh\times2$ ，MLP 隐层约 $2\times BT\cdot 4h\times 2$ （SwiGLU 两个分支）……把这些项按 $BTh$ 归一，就得到「每 token 每 hidden 元素约 34 字节」这个系数，故上式与逐项相加同量级。

与模型状态对比。 若这是一个 7B 模型， $M_{\text{state}}=112$ GB。 $34$ GB 的激活已经与模型状态同一量级，且注意两件事：

它随 $T$ 线性、随 $B$ 线性。把 $T$ 从 2048 提到 8192、 $B_{\text{micro}}$ 从 4 提到 16（等效批放大 16 倍），激活涨到约 $34\times16\approx544$ GB —— 于是即使模型状态已经分片到 14 GB/卡，激活仍然会把卡撑爆；
若注意力矩阵被实体化， $B\,n_h\,T^{2}\times2$ 字节 $=4\times32\times2048^2\times2\approx1.07\times10^{9}$ 字节 ≈ 1 GB 每层，32 层约 34 GB —— 这一项没有被计进上面那个 34 字节系数里（Megatron 的估算把它拆成单独的 $5as/h$ 项），FlashAttention 正是为省掉它而生；它单独就能把预算吃光。

结论：激活账的调参旋钮是 $T$ （序列长度）、 $B_{\text{micro}}$ （单次 token 数）与「是否实体化注意力」，而模型状态账的旋钮是 $N$ 、精度与分片度。 $\blacksquare$

∎

9.1.4 三堵墙：显存、算力、通信

把问题当成「显存不够」会漏掉另外两堵墙。三者的定义与判据如下。

定义 9.1.2 · 大规模训练的三堵墙

设单次训练步骤需要前向计算量 $C_{\text{fwd}}$ （FLOPs 量级，稠密 Transformer 常用近似 $C \approx 6NT$ ，即前向 $2NT$ 、反向 $4NT$ ）、每卡峰值算力 $F$ （FLOP/s）、每卡显存 $M_{\text{cap}}$ 、单步通信量 $V$ （字节）、有效带宽 $\beta$ （B/s）。

墙	判据	症状	处方
显存墙	$M_{\text{state}}+M_{\text{act}}+M_{\text{comm}} \gt M_{\text{cap}}$	OOM、被迫缩小批 / 序列	分片（9.2）、切层（9.3、9.4）、省（9.5）
算力墙	单步时间 $t \gg C/F$ 且 GPU 利用率高	训练慢、但显存没满	提高精度吞吐（FP8）、更大批、更少冗余计算
通信墙	$t_{\text{comm}} \approx V/\beta$ 与 $t_{\text{compute}} \approx C/F$ 可比	GPU 利用率忽高忽低、扩展效率随卡数下降	通信重叠、换拓扑、降并行度、梯度压缩（9.6）

三堵墙的联动关系

它们不是独立的：为了绕开显存墙而切分，切分立刻制造通信墙。所以「加卡」这件事同时作用在两个方向——

加卡缓解显存墙（每卡分摊更少）与算力墙（总算力更大）；
加卡加剧通信墙（更多卡要同步，且跨机链路更慢）。

于是最优卡数既不是 1 也不是「越多越好」，而是「算力墙还紧、通信墙还没到」的那个区间。这也是为什么并行训练的经验总是落在「扩展效率」而不是「总算力」上。

定义 9.1.3 · 加速比与扩展效率

设单卡跑一个训练步骤耗时 $t_1$ ，用 $k$ 张卡跑同样的全局批耗时 $t_k$ （保持全局批不变，即每卡批缩小 $k$ 倍；若保持每卡批不变则可改用强/弱扩展的另一种口径，结论方向相同）。定义

$\text{加速比}\ S=\frac{t_1}{t_k},\qquad \text{扩展效率}\ E=\frac{S}{k}=\frac{t_1}{k\,t_k}.$

$E=1$ 是理想线性加速； $E=0.5$ 意味着一半的卡在等通信或空转。

扩展效率是唯一的验收指标

一个常见误判：卡数翻倍、总吞吐也翻倍 → 以为成功了。但如果每卡吞吐减半，那么总吞吐等于没变，而成本翻倍。所以验收必须看 $E$ ，并且用同样的全局批比较。

量级参考（社区经验值，仅作起点，必须用你自己的集群实测校准）：张量并行留在节点内、流水线并行中 micro-batch 数取到 $M\ge 4P$ 、数据并行配合足够大的全局批， $E$ 通常在 0.7–0.9；跨机通信占比高的配置会掉到 0.5 以下。

例 9.1.3 · 从两段实测反推瓶颈在哪

某训练任务：1 卡单步 $t_1=2.0$ s；8 卡（保持全局批不变）单步 $t_8=3.2$ s。求加速比与扩展效率；并判断瓶颈更可能在通信还是显存。

解

$S=\frac{2.0}{3.2}=0.625,\qquad E=\frac{S}{8}=0.078 .$

$E\approx 7.8\%$ ：8 卡只换来 $0.625$ 倍加速——比单卡还慢。这不是「效率不高」，而是配置本身有问题。常见解释有两种，需要进一步测量区分：

通信占主导：单步通信量 $V$ 相对计算量 $C$ 过大（如模型太小而卡数太多，9.2.2 的通信量公式与 $N$ 无关地摊在很少的计算上）。判据：torch.profiler 里通信内核时间占比高。
配置错误：每卡批太小导致算力利用率低（矩阵乘退化成向量运算），或流水线气泡过大（9.4.2）。

正确的下一步不是「再加卡」，而是先把 $E$ 修复到 0.5 以上：换成更少的卡、更大的全局批、更好的拓扑。 $\blacksquare$

∎

9.1.5 「放不下」的三种情形与决策树

把前面三小节的结论合起来，就能把模糊的「放不下」拆成三种可诊断的情形。

三种「放不下」与对应处方

情形	判据	表现	首选处方	为什么不是别的
A. 模型状态放不下	$16N \gt M_{\text{cap}}$	一启动就在 `optimizer.step()` 后 OOM；模型甚至加载不进去	ZeRO-2/3、FSDP2（9.2）；必要时叠加 8-bit 优化器（9.5.4）	分片不动模型结构、不引入新层内通信；张量 / 流水线并行的工程量更大
B. 激活放不下	$M_{\text{state}}$ 已能放下，但加大批 / 序列就 OOM	前向能跑，反向时 OOM（峰值出现在第一层的反向）	梯度累积 + 激活重计算（9.5.3）；长序列再叠序列并行（9.3.4）	分片对激活无效（9.1.3）；减批会改变优化行为（可用梯度累积补偿）
C. 单层放不下	单个权重矩阵或激活就超过 $M_{\text{cap}}$	某一层（大词表 embedding、超宽 FFN、超长序列注意力）单独就报 OOM	张量并行（9.3）	ZeRO / 流水线都以「层」为粒度，层内放不下时它们无能为力

诊断顺序：先算 $16N$ （9.1.2）排除 A → 再看 OOM 出现在前向还是反向（反向 = B） → 最后单测某一层是否单独就能撑爆显存（= C）。

为什么诊断顺序不能颠倒

因为三条处方的代价递增：

A 类处方（分片）只需改框架配置，几行代码；
B 类处方（梯度累积 + 重计算）要改训练循环，并且改变数值行为（等效批、学习率）；
C 类处方（张量并行）要改模型代码的每一层实现，还引入最频繁的通信。

先做便宜的诊断，能避免「为了一个可以靠配置解决的问题去重写模型」。工程上最常见的浪费就是：明明是 A（优化器状态装不下），却先去做 C（手写张量并行）。

图注：本讲的决策主线（本教程绘制）。三种「放不下」在左中右三条支路上各有处方，但它们最终都会撞上同一堵墙——切分带来的通信。因此任何方案都要回到 [定义 9.1.3](#def-9.1.3) 的扩展效率上验收，而不是以「显存够用了」为终点。

9.1.6 代码：先把账实测出来

估算只是起点，实测对账才是标准动作。下面这段片段在任意一次训练里都能跑，用来找出真实峰值出现在哪一项。

import torch

def report_memory(tag: str, device=0):
    torch.cuda.synchronize(device)
    alloc = torch.cuda.memory_allocated(device) / 1e9        # 当前活跃张量
    peak = torch.cuda.max_memory_allocated(device) / 1e9     # 到目前为止的峰值
    rsvd = torch.cuda.memory_reserved(device) / 1e9          # 缓存池保留量（含碎片）
    print(f"[{tag}] allocated={alloc:.2f}GB peak={peak:.2f}GB reserved={rsvd:.2f}GB")

# ① 模型状态：加载后、第一次前向之前
report_memory("after-model-init")        # 应与 2N 字节量级对得上（bf16 参数）
report_memory("after-optimizer-init")    # 增量应接近 12N（m,v 与主权重 4N+8N）

# ② 单步：分别在三个时刻各测一次，定位峰值出现在哪个阶段
out = model(batch)                       # 前向
report_memory("after-forward")           # 增量 ≈ 激活量（反向结束前不能释放）
loss = out.loss
loss.backward()                          # 反向：梯度开始占用 2N
report_memory("after-backward")          # 峰值通常出现在反向的早期层
optimizer.step(); optimizer.zero_grad(set_to_none=True)
report_memory("after-step")

# ③ 与手算对账：把实测值与 16N / 激活估算并列写进实验记录

与手算对账的三条经验规则

allocated 对账账本、peak 对账预算：判断「装不装得下」看 peak。
reserved - allocated 是碎片与缓存：它常常是那个「明明只用了 60 GB 却 OOM」的原因。通过 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True 一类设置缓解（框架与版本不同，开关名需以你的版本为准）。
拆开测：after-model-init、after-optimizer-init 两个数就能验证「模型状态 = 16N」是否成立；若对不上，先查是否漏了主权重或用了 fp32 梯度。

本节六个高频误解

「显存 = 参数量 × 4 字节」。fp32 参数只是账本的一项。混合精度 Adam 下静态账是 $16N$ ，是「参数 fp32」的四倍。
「减小批大小能让优化器状态变小」。批大小只影响激活；模型状态与批大小无关。反过来，加大批大小也不会让模型状态变大。
「开了混合精度就省一半显存」。省的是参数与梯度这两项：从 $4N+4N$ 降到 $2N+2N$ ；但低精度训练必须另存 fp32 主权重 $4N$ ，正好把这 $4N$ 省下的额度吃回去。所以三种写法的静态账大致是：纯 fp32 Adam $=4N+4N+8N=16N$ ；混合精度 Adam $=2N+2N+4N+8N=16N$ ；混合精度 + 8-bit 优化器 $=10N$ 。结论：混合精度的显存收益主要体现在激活与临时张量上，静态账本几乎不变；它的主要红利是算力吞吐（16 位张量核心）。指望它单独救 7B 会落空（详见 9.5.1 的对照表）。
「OOM 就是模型太大，必须切模型」。先分清三种情形。如果是 B（激活），切模型毫无帮助。
「扩展效率 50% 说明硬件不行」。 $E$ 低多数时候是配置问题：批太小、并行度过高、通信未重叠、拓扑不对。先改配置再看硬件。
「账算得越细越准」。账本是下界：通信缓冲、CUDA context（几百 MB 到数 GB）、碎片、框架临时张量都不在四项里。工程上按算出的量级再留 15%–30% 余量是常见做法。

即时自测 · 9.1

Q1 用一句话说清「模型状态」与「激活」的区别，并各举一种只对它有效的省显存手段。

Q2 某模型 $N=1.5\times10^{9}$ ，混合精度 Adam 训练。手算模型状态显存；一张 24 GB 卡在不分片、无其他优化的前提下能否放下？（写出算式）

Q3 训练到 backward() 时 OOM，但模型加载与 optimizer.step() 之后都正常。按三种「放不下」与对应处方的分类，这属于哪种情形？给出两条首选处方。

Q4 8 卡训练实测 $t_1=2.0$ s、 $t_8=3.2$ s（全局批不变）。至少给出两个「不该做」的动作，并说明为什么。

参考要点

Q1 模型状态 = 参数 + 梯度 + 优化器状态 + 主权重，与 $N$ 成正比、与批大小无关、训练全程常驻（处方：ZeRO / FSDP 分片、8-bit 优化器、降低精度）；激活 = 前向中间张量，与「单次 token 数 × 层数」成正比、用完即可释放（处方：梯度累积 / 减小 micro-batch、激活重计算）。两者的手段不可互换：分片对激活无效，减小 micro-batch 对模型状态无效。

Q2 $M_{\text{state}}=16N=16\times1.5\times10^{9}=24\times10^{9}$ 字节 $=24$ GB。再加激活与碎片必然超过 24 GB —— 放不下。注意这里还有一个细节：24 GB 卡的实际可用显存小于 24 GB（驱动与 CUDA context 会先占掉一部分），所以即使算出来「刚好 24 GB」也不能用。

Q3 属于 B. 激活放不下（OOM 出现在反向，而模型状态与优化器初始化都已通过）。处方：① 梯度累积把等效批留住，同时把 $B_{\text{micro}}$ 降到能容纳的水平；② 开激活重计算（全层或选择性）。若序列很长再加序列并行。

Q4 不该做：① 再加卡—— $E$ 只有 $7.8\%$ ，加卡只会让通信占比更高；② 直接换更大模型 / 更长序列——当前配置连基本效率都没达到，放大只会掩盖问题。也可答「把每卡批调更小」：那会进一步降低算力利用率。正确方向是减少并行度、增大全局批、检查通信占比与拓扑（9.6.4）。

Keyboard shortcuts

深度学习导论 · 定制教程