推理任务能放下，为什么训练任务放不下#
上一篇讨论了模型调度中的 Capacity Constraint：一组模型到底能不能同时放进一张 GPU。
那篇文章的结论是：推理场景下，判断模型能否部署到 GPU，先看 persistent memory（主要是模型权重），再加上 forward 过程中的 activation、workspace 和运行时开销。最终关心的是某一时刻"同时还活着"的 tensor 有多少——也就是 peak live memory。
但如果把其中一个推理任务换成训练任务，这套估算很快就会失效。
同一个模型，参数没有变，推理时可能只需要几个 GB。一旦开始训练，显存可能膨胀到推理的数倍甚至十几倍。更麻烦的是，训练显存甚至不是一个稳定的数字——它在 forward、backward 和 optimizer step 之间不断变化，峰值出现在哪个阶段都不一定。
这里有一个看起来反直觉的现象：activation checkpointing 会主动丢掉 forward 已经算出来的中间结果，在 backward 时重新计算一遍。“把计算做两遍"怎么反而是一种重要的训练优化？
理解这个问题，需要先理解训练显存到底由什么决定。
从推理显存账本到训练显存账本#
上一篇建立了推理场景的 mental model：
$$\boxed{M_{\text{infer, peak}} \approx M_{\text{persistent}} + M_{\text{peak\ runtime\ working\ set}}}$$训练在这个基础上至少额外引入了：
- Parameter gradients：每个参数对应一份梯度
- Optimizer states：Adam 需要维护一阶矩和二阶矩
- FP32 master weights：某些混合精度实现中保留一份全精度参数副本
- Retained activations：为了 backward 而长期保存的 forward 中间结果
- Communication buffers：分布式训练中的梯度同步和参数聚合
于是训练的显存近似变成：
$$M_{\text{training\ peak}} \approx M_{\text{parameter}} + M_{\text{gradient}} + M_{\text{optimizer}} + M_{\text{activation}} + M_{\text{temporary}} + M_{\text{communication}} + M_{\text{runtime}}$$把这些状态按两种角色分开更清楚。
与模型规模相关的常驻状态#
- Parameters
- Gradients
- Optimizer states
- FP32 master weights（如果有）
这些主要随参数量增长，和具体的输入数据关系不大。
与执行计划相关的动态状态#
- Forward activations
- Operator workspace
- Communication buffers
- Allocator reserved memory
- 同时处于 in-flight 状态的 micro-batches
这些不仅取决于模型本身，还取决于 micro-batch size、sequence length、hidden size、forward/backward 调度、kernel 实现、并行方式、tensor 生命周期。
flowchart LR
subgraph Inference["推理"]
I_P["Parameters"]
I_A["Activations"]
I_W["Workspace"]
I_R["Runtime"]
end
subgraph Training["训练"]
T_P["Parameters"]
T_G["Gradients"]
T_O["Optimizer States"]
T_M["Master Weights"]
T_A["Retained Activations"]
T_W["Workspace"]
T_C["Comm Buffers"]
T_R["Runtime"]
end
Inference -->|"同一个模型"| Training
上一篇讨论过 MoE 的例子：compute 和 memory 是两个独立的资源维度。这里出现了类似的分离：模型权重只描述了模型有多大；训练配置和执行计划决定了这个模型跑起来时有多大。
Backward 为什么需要保留 forward 的中间结果#
先用最简单的线性层看清楚这件事。
Forward：
$$Y = WX$$Backward 中，已知上游传回的 \(\frac{\partial L}{\partial Y}\)，需要计算两个梯度：
$$\frac{\partial L}{\partial W} = \frac{\partial L}{\partial Y} X^T$$$$\frac{\partial L}{\partial X} = W^T \frac{\partial L}{\partial Y}$$注意第一个式子：计算 parameter gradient \(\frac{\partial L}{\partial W}\) 需要 forward 阶段的输入 \(X\)。
这意味着：\(X\) 不能在 forward 结束后立即释放。它必须一直活到对应的 backward pass 完成。
这就是训练显存和推理显存最根本的区别。推理时，一个 activation 被下一层使用后通常可以立即释放。训练时，backward 强制延长了 activation 的 lifetime。
而且不只是一层。一个 N 层网络的第一层输入，可能要等到 backward 一路回传到第一层时才能释放——也就是整个 forward 加上几乎整个 backward 期间都得常驻。
需要明确一个术语：训练系统语境中的 activation，不只是激活函数的输出，而是所有为了 backward 而保留的前向中间结果。 Autograd 框架会根据 computation graph 的依赖关系，自动决定哪些中间 tensor 需要保存。
参数、梯度和 Optimizer State 的数量级#
先看数量级差异是怎么来的。以 BF16 mixed-precision Adam 为例，粗略账本：
| Component | bytes/parameter |
|---|---|
| BF16 parameter | 2 |
| Gradient（BF16 或 FP32） | 2~4 |
| Adam 一阶矩（FP32） | 4 |
| Adam 二阶矩（FP32） | 4 |
| FP32 master weight（如果有） | +4 |
因此经常看到 12~18 bytes/parameter 这样的粗估。工程上也常用大约 16 bytes/parameter 做第一轮容量估算。
对比一下：
- BF16 推理：主要是约 2 bytes/parameter 的模型权重
- BF16 mixed-precision Adam 训练：仅模型状态就可能接近 16 bytes/parameter
同一个模型，参数量没变，但训练任务的显存账本已经是推理的 8 倍。而且这还没算 activation、workspace 和通信 buffer。
必须强调：12~16 bytes/parameter 只是某些 mixed-precision Adam 配置下的估算，不是定律。换 optimizer（SGD、Adafactor、LAMB）、换 precision scheme（纯 BF16、FP8）、加 ZeRO/FSDP sharding，单卡 footprint 都会变。关于不同 dtype 的 bytes/element 和 mixed precision 训练中 storage dtype 与 compute dtype 的区别，见 IEEE 754 到 BF16 和 Mixed Precision 两篇。
显存不是一个数字，而是一条时间线#
上面的 bytes/parameter 估算描述的是"这些状态最终都需要多少空间”。但它没有回答一个更关键的问题：这些状态什么时候同时存在？
一次训练 iteration 中，显存占用是随执行阶段变化的：
- Parameters 几乎全程常驻
- Forward 阶段：activations 不断积累，越来越多的中间结果被保留等待 backward
- Backward 阶段：gradients 逐层产生，同时对应层的 saved activations 被使用后逐渐释放
- Optimizer step：parameters、gradients 和 optimizer states 同时参与更新，可能产生新的临时峰值
- 首次 optimizer step 还可能延迟创建 optimizer states（lazy initialization），在这一刻突然分配大量显存
gantt
title 训练 Iteration 的显存生命周期
dateFormat X
axisFormat %s
section 常驻状态
Parameters :p, 0, 30
OptimizerStates :os, 0, 30
section Forward
Act_L1 :a1, 2, 22
Act_L2 :a2, 4, 20
Act_L3 :a3, 6, 18
Act_L4 :a4, 8, 16
section Backward
Grad_L4 :g4, 14, 24
Grad_L3 :g3, 16, 24
Grad_L2 :g2, 18, 24
Grad_L1 :g1, 20, 24
section Optimizer
OptimizerStep :opt, 24, 28
观察这条时间线，有几件事变得很清楚：
Forward 和 backward 交界处附近，大量 activations 和刚开始产生的 gradients 同时存在——这通常是整个 iteration 中显存占用最高的区域。
而 OOM 看的正是时间轴上的峰值，不是平均显存，也不是模型加载后的静态占用。
关于 allocated、reserved 和 nvidia-smi 之间的区别，上一篇和 Caching Allocator 那篇已经讨论过，这里不重复。核心是同一个：tensor 逻辑上释放不等于物理显存立即归还。
一个简短的诊断框架#
当训练 OOM 时，主要矛盾出现在什么地方？
- 模型或 optimizer 初始化时就放不下：优先看 parameters 和 optimizer states 的总量
- batch 或 sequence length 增大后 OOM：优先看 activations
- 某个算子执行时突然达到峰值：检查中间 tensor 和 workspace
- 多卡总显存很多，但单个任务仍放不下：检查当前并行方式是在复制还是分片
nvidia-smi与框架统计差异很大：检查 reserved memory、CUDA context 和外部库分配
Activation Checkpointing 改变了什么#
现在可以回到开头那个反直觉的问题：为什么"把计算做两遍"反而是一种优化？
正常训练中：
- Forward 生成中间结果
- 所有需要 backward 使用的 activation 一直保留
- Backward 逐层使用后才释放
上面的时间线图已经展示了问题：早期层的 activation（如 Act_L1）从 forward 一开始就产生，一直要等到 backward 几乎结束才能释放。存活时间非常长。
Activation checkpointing 的做法是：
- 只保留少数 边界 activation（checkpoint 点）
- 丢弃 checkpoint 区间内部的中间结果
- Backward 到达该区间时，从边界 activation 出发，重新执行部分 forward
- 使用重新生成的中间结果完成 backward
- 随后立即释放
仅仅说"以计算换显存"还不够精确。Activation checkpointing 没有改变模型参数——参数量不变、单个 tensor 不变。它改变的是：
- 哪些 tensor 需要长期保存
- Tensor 何时重新生成
- 同一时刻有多少 tensor 同时存活
Activation checkpointing 真正缩短的是中间 Tensor 的存活时间。
从 FLOPs 角度看（上一篇讨论过 forward 约 \(F\)、backward 约 \(2F\)、正常训练约 \(3F\)）：极端情况下完整重算一次 forward，总计算量约为 \(4F\)，增加约三分之一。但实际开销取决于 checkpoint 粒度、kernel 效率和硬件利用率。
同时需要认清 checkpointing 的边界：
- 主要减少 retained activations
- 不减少 parameters
- 不减少 optimizer states
- 不一定降低单个算子执行时的 workspace 峰值
因此，如果仅模型状态（parameters + optimizer states）就已经放不下一张 GPU，activation checkpointing 无法解决根本问题。
训练显存优化的几种基本手段#
改变表示：让 Tensor 更小#
代表：BF16/FP16、FP8、低精度 optimizer states、8-bit optimizer。
本质上在做：
$$\text{显存和带宽} \quad\Longleftrightarrow\quad \text{数值精度、稳定性和实现复杂度}$$关于不同 dtype 的 trade-off，IEEE 754 到 BF16 已经讨论了一个 number 内部如何分配 bit budget，Mixed Precision 讨论了整个 graph 怎么分配 precision budget。这里不重复。
需要说明的是：低精度不意味着全部训练显存自动减半。因为 optimizer states（Adam m/v）通常仍然是 FP32，master weights 可能保留 FP32，部分数值敏感的算子也会保持高精度。Mixed precision 是 precision allocation strategy，不是把所有东西对半砍。
改变生命周期：让 Tensor 少同时存在一会儿#
代表：Activation checkpointing、Pipeline Parallel 中的 1F1B 调度。
上面已经分析了 activation checkpointing。1F1B（one forward one backward）的思路类似：通过交替执行不同 micro-batch 的 forward 和 backward，控制同一时刻有多少 micro-batch 的 activation 同时存在，而不是等所有 micro-batch 的 forward 都做完再统一 backward。
核心 insight：
调度器不仅关心一块内存有多大，还关心它从什么时候占用到什么时候。 训练显存不是若干静态数字之和，而是一组存在重叠关系的生命周期。这和 Caching Allocator 那篇讨论的 allocation lifetime 是同一个问题。
改变物化方式：不要创建完整中间结果#
代表：FlashAttention、fused cross entropy、算子融合、分块和在线计算。
需要区分两种"减少中间结果"的方式：
- Checkpointing：中间结果在 forward 时确实产生过，但暂时不保存，backward 时重新计算
- FlashAttention：改变计算顺序，让完整的 \(N \times N\) attention 矩阵从一开始就不需要被写入 HBM
比压缩或搬运中间结果更彻底的优化，是让它从一开始就不要存在。
而且减少中间 tensor 往往也减少 HBM 流量。正如 Arithmetic Intensity 那篇讨论的：memory-bound 操作的瓶颈在数据搬运。减少中间 tensor 的物化，可能同时节省显存和提升性能。
改变并发粒度：减少同一时刻处理的数据#
代表：减小 micro-batch、gradient accumulation、调整 pipeline 中的 micro-batch 数量。
Gradient accumulation 把一个大 batch 拆成多次小规模 forward/backward，在多轮小 batch 上累加梯度，然后一次 optimizer step。
决定 activation 峰值的是 micro-batch size，而不是最终的 global batch size。
Gradient accumulation 没有消除计算，也没有消除 parameter gradients。它用更长的执行时间换取更低的单次峰值。可能付出的代价包括：GEMM shape 变小导致 GPU 利用率下降、iteration 时间增加、同步和训练语义更加复杂。
改变放置位置：把状态分散到其他 GPU#
代表：DDP、ZeRO/FSDP、Tensor Parallel、Pipeline Parallel、Sequence/Context Parallel。
不要详细推导并行算法，只看它们在切分什么：
| 方法 | 主要切分什么 | 主要解决什么 | 主要代价 |
|---|---|---|---|
| DDP | 数据和计算 | 提高吞吐 | 模型状态仍被复制 |
| ZeRO/FSDP | 参数、梯度、optimizer states | 单卡模型状态 | 通信和参数聚合 |
| Tensor Parallel | 单层权重与计算 | 单层过大 | 高频 collective |
| Pipeline Parallel | 模型层 | 模型过深 | bubble 与调度 |
| Sequence/Context Parallel | 序列维度 | 长序列 activation | token 间通信 |
这里有一个非常值得注意的区别：
普通 DDP 增加的是总算力和总显存，却不一定增加单个任务可使用的有效显存。 因为每张 GPU 仍然复制完整的模型状态。
ZeRO/FSDP 的核心变化是：
$$\text{长期复制} \longrightarrow \text{分片保存、按需聚合}$$它没有让模型状态消失，而是用通信换取单卡容量。每张卡只持有参数、梯度和 optimizer states 的一个分片，需要时通过 all-gather 聚合完整参数。
改变存储层级：把状态搬到更便宜的内存#
代表：CPU offload、NVMe offload、optimizer state offload、activation offload。
核心是：
$$\text{GPU HBM 容量} \quad\Longleftrightarrow\quad \text{PCIe、内存或存储带宽与延迟}$$Optimizer states 在每轮 iteration 中只使用一次（optimizer step），使用频率相对较低，因此更适合 offload 到 CPU 内存。
Activation 在 backward 中很快就会再次被读取，如果 offload 后读不回来够快，系统会受传输带宽限制。
Offload 不是"免费增加显存"，而是把容量问题转化成数据移动问题。
训练显存优化的统一视角#
上面六种手段看起来各不相同，但可以用同一组问题归纳：
- 这个 Tensor 是否真的需要被完整创建？
- 如果需要，它能不能使用更低精度？
- 它必须保存多久？
- 同一时间必须有多少份数据存在？
- 它一定要放在这张 GPU 上吗？
- 为了减少显存，需要增加多少计算、通信、传输或系统复杂度？
| 优化方式 | 改变了什么 | 用什么换显存 |
|---|---|---|
| Mixed precision | Tensor 的表示 | 数值精度与稳定性 |
| Activation checkpointing | Tensor 的生命周期 | 重计算 |
| Gradient accumulation | 同时处理的数据量 | 吞吐和更多迭代 |
| FlashAttention / 算子融合 | 中间结果的物化方式 | Kernel 和算法复杂度 |
| ZeRO / FSDP | 状态所在的设备 | 网络通信 |
| Pipeline / Tensor Parallel | 模型的切分方式 | 通信与调度 |
| CPU / NVMe offload | 状态所在的内存层级 | 带宽与延迟 |
每一种显存优化，都在把压力从 HBM 转移到另一种资源：计算、通信、带宽、精度或系统复杂度。真正的工程问题，是在当前系统中哪一种资源更便宜。
从显存优化回到模型调度#
讲完这些优化后，必须回到模型调度，而不是停留在训练算法层面。
推理场景下，上一篇已经指出 model.memory 是一个危险的抽象。到了训练场景，这个问题更加严重。
一个调度器判断训练任务能否放到某张 GPU 上，不能只记录：
$$\text{model\_size} = X \text{ GB}$$它至少需要理解：
- 模型状态的常驻占用（parameters + gradients + optimizer states）
- micro-batch 和序列长度决定的 activation 峰值
- forward、backward、optimizer step 各阶段的显存峰值
- activation checkpointing 是否开启
- optimizer 和混合精度策略
- 当前使用 DDP、FSDP 还是其他并行方式
- 通信和参数聚合的临时 buffer
- allocator reserved memory、CUDA context 和安全余量
训练任务的资源需求不是一个静态的"模型大小"，而是一条由执行计划决定的显存曲线。
Activation checkpointing、gradient accumulation、1F1B 等方法，本质上都在改变这条曲线的形状。Mixed precision 改变曲线的高度；checkpointing 改变 activation 的重叠范围；FSDP 改变状态在哪些设备上存在；offload 改变状态存放在哪一级内存。
训练显存优化不仅决定任务能不能跑，也直接改变任务应该如何被调度。
结论#
在推理场景中，我们容易把模型显存理解成"权重加上一些运行时开销"。上一篇建立的 mental model 是 Persistent Memory + Peak Runtime Working Set，对推理已经够用。
到了训练阶段，这种静态视角就不够了。
训练引入的不只是 gradients 和 optimizer states——这些只是让显存账单变大。更重要的变化是：backward 彻底改变了 tensor 的生命周期。早期层的 activation 必须存活到 backward 几乎结束，大量中间结果在 forward-backward 交界处同时存在。
Activation checkpointing 之所以重要，不只是因为它能省下一部分 activation。它说明了一件事：已经计算出来的中间结果，并不一定要一直保存。
从这个角度看，训练显存不再只是一个容量问题，而是一个调度问题：
- 哪些状态常驻
- 哪些状态按需生成
- 哪些状态可以分片
- 哪些状态可以搬运
- 哪些中间结果根本不应该被物化
模型大小只是模型本身的属性；训练显存是执行计划的属性。
显存优化的本质，是重新安排 Tensor 在何时、何地、以什么形式存在。
参考资料#
- 从参数量到 Peak Memory：模型显存到底应该怎么算？（本文前篇）
- 从 FLOPs 到 Latency：GPU 推理性能到底由什么决定？（姊妹篇）
- 从 IEEE 754 到 BF16：理解 ML Infra 中的浮点精度选择（姊妹篇）
- 从 AMP 到 Mixed Precision：训练和推理为什么不使用同一种精度（姊妹篇）
- 从线上模型切换问题理解 PyTorch CUDA Caching Allocator（姊妹篇）
- 从 CS336 到 ML Infra：理解 Arithmetic Intensity（姊妹篇）
- Rajbhandari et al. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. SC 2020.
- Chen et al. Training Deep Nets with Sublinear Memory Cost. 2016.
- Dao et al. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS, 2022.
- Korthikanti et al. Reducing Activation Recomputation in Large Transformer Models. MLSys, 2023.