4. 计算图与显存：动态图与静态图的显存差异、中间激活值的存储、反向传播的显存开销
这一章，我们来聊聊计算图。说实话，很多同学刚开始调显存的时候，根本没意识到计算图才是真正的“幕后黑手”。你想想看，同样的模型，换个框架跑，显存占用能差出一倍。为什么？说白了，就是动态图和静态图在背后搞的鬼。
4.1 动态图 vs 静态图：显存管理的两种哲学
我个人习惯把动态图叫做“即用即弃”模式。PyTorch 就是典型代表。你每执行一行代码，它就当场构建一个节点，前向传播完了，计算图就扔在一边，等着反向传播时再用。
静态图则完全不同。TensorFlow 1.x 或者 JAX 那种方式，是先搭好一个完整的“施工蓝图”，然后一次性执行。图一旦建好，中间哪些节点该保留、哪些可以扔掉，编译器心里门儿清。
核心差异一句话：动态图为了灵活性，不得不保留完整的计算图直到反向传播结束；静态图因为提前知道了全局结构，可以更激进地回收显存。
我在项目中遇到过一件事。有一次用 PyTorch 跑一个深层 ResNet，batch size 设到 32 就 OOM 了。换成 TensorFlow 2.x 的 @tf.function 静态图模式，同样的硬件，batch size 直接拉到 48 还没爆。嗯，这就是静态图的“预知能力”带来的红利。
4.2 中间激活值的存储：显存消耗的“大头”
很多人以为模型参数占显存最多。其实错了。真正的大头，是中间激活值（Intermediate Activations）。
什么叫中间激活值？就是每一层算出来的那个输出，还没经过下一层之前，得先找个地方放着。反向传播的时候，需要用它来算梯度。你想想看，一个 100 层的网络，每一层都要存一份完整的 feature map，这显存能不大吗？
| 组件 | 显存占比（典型情况） | 说明 |
|---|---|---|
| 模型参数 | 10% - 20% | 权重和偏置，相对固定 |
| 中间激活值 | 50% - 70% | 每层的输出，反向传播必需 |
| 梯度 | 10% - 15% | 反向传播时临时产生 |
| 优化器状态 | 10% - 20% | 如 Adam 的动量、方差 |
为什么会这样？因为反向传播的计算公式决定了：
# 假设一个简单的全连接层
# 前向：z = W @ x + b
# 反向：dL/dW = dL/dz @ x.T
# 看到了吗？反向传播需要用到前向的输入 x
# 所以 x 必须被保留下来，直到反向传播算完
我曾经踩过一个坑。训练一个 Transformer 模型，序列长度从 512 改成 1024，显存直接翻了三倍。一开始我还以为是参数变多了，后来一查才发现，是中间激活值随着序列长度平方级增长。嗯，自注意力机制的 QK^T 矩阵，那个尺寸是 L×L 的，存下来可不是闹着玩的。
我的建议：如果你在做长序列任务，优先考虑激活值检查点（Activation Checkpointing）。它用计算换显存，只保留部分层的激活值，其他层在反向传播时重新算一遍。虽然慢了 20%-30%，但显存能省下一大半。
4.3 反向传播的显存开销：为什么反向比前向更“吃”显存？
很多人以为反向传播只是算算梯度，不占什么显存。其实恰恰相反。反向传播的显存开销，往往比前向传播还要大。
原因有三：
- 需要保留前向的“证据”：每一层的输入、输出，都得留着。反向传播时才能算出对权重的梯度。
- 梯度本身也占地方：每个参数都有一个对应的梯度，尺寸和参数一模一样。
- 中间梯度链：链式法则要求我们一层一层往回传，每一层的中间梯度都得暂存。
举个例子，一个三层的网络：
# 前向传播
x1 = layer1(x0) # 存 x0, x1
x2 = layer2(x1) # 存 x1, x2
x3 = layer3(x2) # 存 x2, x3
loss = loss_fn(x3, y) # 存 x3
# 反向传播
# 需要用到 x3 算 dL/dx3
# 需要用到 x2 算 dL/dx2
# 需要用到 x1 算 dL/dx1
# 需要用到 x0 算 dL/dx0
# 直到所有梯度算完，这些中间值才能释放
你看，前向传播只需要存一份当前层的输出。但反向传播需要同时持有多个层的中间值。这就是为什么反向传播的显存峰值往往出现在反向刚开始的时候——所有需要的东西都堆在那里。
避坑指南：我曾经在训练一个 GAN 的时候，生成器和判别器交替更新，显存突然暴涨。查了半天才发现，是因为 PyTorch 默认会保留所有参与计算的张量的计算图。如果你在循环里没有显式调用 .detach() 或者 .requires_grad_(False)，计算图会越积越大，最后直接 OOM。
4.4 知识体系：计算图与显存的核心逻辑
下面这张图，是我自己总结的计算图显存管理的核心逻辑。你看一遍，应该就能把前面讲的东西串起来了。
这张图想表达的核心思想是：动态图和静态图在最上层就分道扬镳了，但它们最终都要面对中间激活值的存储问题。而反向传播的显存峰值，本质上是由“必须保留前向证据”这个硬约束决定的。优化策略，说白了就是在“保留什么”和“重新算什么”之间做权衡。
一个小技巧：如果你用 PyTorch，可以试试 torch.cuda.memory_summary() 这个函数。它能告诉你每一块显存是干什么用的——哪些是参数，哪些是激活值，哪些是梯度。我每次调优之前，都会先跑一下这个，看看“钱”到底花在了哪里。
好了，这一章的内容就到这里。计算图这个东西，你越早理解它，后面调显存的时候就越有底气。记住一句话：显存优化的本质，就是管理好计算图中的每一个节点，让它们在该出现的时候出现，在该消失的时候消失。