4、激活值压缩与重计算：激活值存储开销、Checkpointing策略、选择性重计算、内存与计算权衡

4.1 激活值存储开销——被忽视的显存大户

很多人以为训练大模型，显存大头是模型参数。其实不然。我刚开始做大模型部署时，也被这个认知坑过。

你想想看，一个 7B 的模型，参数用 FP16 存，也就 14GB 左右。但训练时呢？激活值（Activation）才是真正的显存杀手。举个例子，一个 batch size 为 4、序列长度 2048 的 LLaMA-7B，单层激活值就能吃掉 2-3GB。几十层堆下来，轻松破百 GB。

为什么会这样？因为前向传播时，每一层的输出都得留着。反向传播要用这些中间结果算梯度。说白了，这就是用空间换时间——不存的话，就得重新算一遍。

我在项目中遇到过最夸张的情况：一个 13B 模型，训练时激活值占了总显存的 70% 以上。参数反而只占 20% 不到。嗯，这时候你就知道，必须对激活值下手了。

4.2 Checkpointing 策略——经典的重计算方案

激活值存太多，那就少存点呗。Checkpointing（也叫梯度检查点）就是这个思路。

它的做法很简单：在前向传播时，只保存部分层的激活值，而不是全部。反向传播时，遇到没存的那层，就重新算一遍。

我习惯把模型分成若干段，每段只存第一个层的输入。这样显存占用能降到原来的 1/N。代价就是多了一次前向计算。

# 伪代码示例：Checkpointing 的核心逻辑
def forward_with_checkpoint(model, x):
    # 只保存 checkpoint 点的激活值
    checkpoints = [x]
    for i, layer in enumerate(model.layers):
        x = layer(x)
        if i % checkpoint_interval == 0:
            checkpoints.append(x)  # 存下来
    return x, checkpoints

def backward_with_checkpoint(grad, model, checkpoints):
    # 从最近的 checkpoint 开始重计算
    for i in reversed(range(len(model.layers))):
        if i % checkpoint_interval == 0:
            # 从 checkpoint 恢复
            x = checkpoints[i // checkpoint_interval]
        else:
            # 重计算前向
            x = recompute_forward(model.layers[i], x)
        # 正常反向传播
        grad = model.layers[i].backward(grad, x)
    return grad

4.3 选择性重计算——别什么都重算

Checkpointing 虽然好，但有个问题：它太粗暴了。所有层一视同仁，该重算的都重算。其实没必要。

我后来发现，不同层的激活值大小差异很大。比如 Attention 层的 QKV 输出，每个 token 都要存，占地方。但 FFN 层的中间激活，其实可以压缩或者重算。

选择性重计算的核心思路就是：只重算那些「计算量小、存储量大」的层。

层类型	激活值大小	重计算开销	建议策略
Attention (QKV)	大 (O(n²d))	中	存储
Attention (Softmax)	极大 (O(n²))	小	重计算
FFN (中间层)	大 (O(nd))	小	重计算
LayerNorm	小 (O(n))	极小	重计算

你看，Softmax 的输出虽然占地方，但重算它只需要做一次 exp 和 sum，成本极低。FFN 的中间激活也是，重算就是一次矩阵乘，比存下来划算多了。

4.4 内存与计算权衡——找到你的甜蜜点

说白了，激活值压缩和重计算，就是在做一道选择题：你要省显存，还是要省时间？

我一般会画一条曲线：横轴是显存节省比例，纵轴是训练时间增加比例。理想情况下，我们希望曲线越陡越好——省一点显存，时间增加很少。

我个人习惯用选择性重计算。它比纯 Checkpointing 多省 5% 显存，但时间只多 10%。性价比最高。

你想想看，如果你的 GPU 刚好差那么点显存，用选择性重计算就能塞下更大的 batch size。batch size 翻倍，训练速度反而能提升 30-40%。这笔账，划算。

上面这张图是我自己总结的。你看，曲线在「选择性重计算」附近最陡。再往左，虽然显存省得更多，但时间成本涨得太快了。再往右，显存省得不够多，意义不大。

所以我的建议是：先做选择性重计算，把 Attention 的 Softmax 和 FFN 中间层重算了。如果显存还不够，再考虑加 Checkpointing。

嗯，关于激活值压缩和重计算，核心就这些。记住一句话：别盲目重算，也别全存。找到你的甜蜜点，才是最优解。

想要更多资料，联系可可