3、显存峰值计算：前向传播显存占用、反向传播显存占用、优化器状态显存占用
聊到Batch调参，大家最头疼的是什么？
我个人觉得，不是模型收敛慢，而是——显存爆了。
你想想看，调大Batch Size，训练速度上去了，结果OOM（Out of Memory）直接给你报错。调小吧，又嫌效率低。这中间的平衡点，说白了就是显存峰值在作祟。
今天我们就来拆解一下，显存到底被谁吃了。我把它分成三块：前向传播、反向传播、优化器状态。这三块加起来，就是你的显存峰值。
核心公式：
显存峰值 ≈ 前向激活显存 + 反向梯度显存 + 优化器状态显存 + 模型参数显存
注意：前向和反向的显存往往有重叠，实际峰值需要看具体实现。
3.1 前向传播显存占用
前向传播时，显存主要用来存两样东西：模型参数和中间激活值。
模型参数好理解，就是权重和偏置。但中间激活值，很多人会忽略。其实它才是显存大户。
举个例子，你有一个线性层：输入是 [Batch, 1024]，输出是 [Batch, 512]。前向传播时，系统需要保存输入张量 [Batch, 1024]，因为反向传播时要算梯度。这个保存下来的东西，就叫激活值。
我遇到过这样一个项目：一个12层的Transformer，Batch Size设成64。我以为参数才几百MB，显存肯定够。结果一跑，直接OOM。查了半天，发现是中间激活值占了将近6GB。
我的经验：
前向传播的显存占用，大致可以这样估算：
- 每层输出张量大小 × Batch Size × 精度字节数
- 所有层的激活值累加
- 加上输入数据本身
举个例子：一个卷积层输出特征图是 [64, 112, 112, 64]，Batch=32，float32。那么这一层激活值就是 64×112×112×64×4 ≈ 205MB。如果你有20层，光激活值就4GB了。
3.2 反向传播显存占用
反向传播时，显存占用主要来自梯度和临时计算缓冲区。
梯度本身不大，因为每个参数对应一个梯度，大小和参数一样。但反向传播过程中，框架会创建很多临时张量来存储中间计算结果。这些临时张量，有时候比激活值还大。
为什么会这样？
我打个比方。前向传播时，你只需要保存每个层的输出。但反向传播时，你需要从损失函数开始，一层层往回算。每一层都要用到前向保存的激活值，同时还要生成新的中间结果。这些中间结果，就是临时缓冲区。
我记得有一次调优一个ResNet-50模型。我把Batch Size从32调到64，心想显存应该翻倍。结果翻了将近3倍。为什么？因为反向传播的临时缓冲区，随着Batch Size增大，增长是非线性的。
避坑指南：
我曾经踩过一个坑：使用PyTorch的自动混合精度训练（AMP）。前向传播时，FP16的激活值确实省显存。但反向传播时，框架为了精度，会把部分梯度转成FP32计算。结果显存不降反升。
所以，不要盲目相信“低精度=低显存”。一定要实测。
3.3 优化器状态显存占用
优化器状态，很多人会忽略。但它往往是压死骆驼的最后一根稻草。
不同的优化器，占用的显存不一样：
| 优化器 | 每个参数额外占用 | 说明 |
|---|---|---|
| SGD | 0字节 | 没有动量，几乎不占额外显存 |
| SGD + Momentum | 4字节（float32） | 保存动量向量，大小和参数一样 |
| Adam | 8字节（float32 × 2） | 保存一阶动量和二阶动量 |
| AdamW | 8字节（float32 × 2） | 和Adam一样，但权重衰减方式不同 |
你想想看，一个1亿参数的模型，用Adam优化器。光是优化器状态就要占：1亿 × 8字节 = 800MB。加上模型参数本身（400MB），前向激活值（假设2GB），反向临时缓冲区（假设1GB）。总显存轻松超过4GB。
我个人的习惯是：先算优化器状态，再算模型参数，最后才是激活值。因为前两个是固定的，激活值还可以通过梯度检查点（Gradient Checkpointing）来压缩。
实战建议：
- 如果显存紧张，优先考虑SGD + Momentum，而不是Adam
- 如果必须用Adam，试试8-bit Adam（bitsandbytes库）
- 优化器状态可以offload到CPU，但会拖慢训练速度
3.4 三块显存的协同关系
这三块显存不是独立存在的。它们之间会互相影响。
举个例子：你减小Batch Size，前向激活值变小，反向临时缓冲区也变小。但优化器状态不变。所以，当Batch Size小到一定程度时，优化器状态就成了显存瓶颈。
反过来，你增大Batch Size，前向激活值和反向缓冲区都变大。但优化器状态不变。这时候，激活值就成了瓶颈。
我画了一张图，帮你理解这三块显存的关系：
嗯，这里要注意：显存峰值不是简单的加法。因为前向传播的激活值，在反向传播完成后就可以释放。所以实际峰值，往往出现在反向传播刚开始的时候——此时前向激活值还在，反向临时缓冲区刚创建，优化器状态也占着。
我的调试技巧：
如果你想知道模型到底吃了多少显存，可以用PyTorch的 torch.cuda.max_memory_allocated() 来监控。我一般会在每个epoch开始和结束时打印一次，看看峰值出现在哪里。
另外，nvidia-smi 显示的显存使用量，往往比实际峰值低。因为它显示的是当前时刻的占用，不是历史峰值。
3.5 如何快速估算显存峰值
我总结了一个快速估算公式，你可以直接套用：
显存峰值 ≈ 参数显存 × (1 + 优化器系数) + 激活显存 × Batch Size + 临时缓冲区
其中：
- 参数显存 = 参数量 × 4字节（float32）
- 优化器系数：SGD=0, SGD+Momentum=1, Adam=2
- 激活显存：每层输出大小 × 4字节
- 临时缓冲区：约等于激活显存的 0.5~1倍举个例子：一个1亿参数的模型，用Adam优化器，Batch Size=64，每层激活值平均50MB，共20层。
参数显存 = 1亿 × 4 = 400MB
优化器状态 = 400MB × 2 = 800MB
激活显存 = 50MB × 20 = 1000MB
临时缓冲区 ≈ 1000MB × 0.8 = 800MB
显存峰值 ≈ 400 + 800 + 1000 + 800 = 3000MB ≈ 3GB这个估算值，和实际跑出来的结果，误差一般在10%以内。我个人用了很久，挺准的。
注意：
这个公式没有考虑框架本身的显存开销（比如CUDA上下文、数据加载器等）。实际使用时，建议留出20%的余量。
我曾经因为没留余量，在训练到第100个epoch时突然OOM。查了半天，发现是数据加载器缓存了太多数据。从那以后，我每次都会多留1GB的buffer。
好了，显存峰值计算这块，核心就是这三块。你只要把前向、反向、优化器状态算清楚，调Batch Size的时候心里就有底了。
想要更多资料，联系可可