4. 重计算（Recomputation）技术：激活值重计算原理、选择性重计算策略、重计算与流水线并行的结合、重计算带来的计算开销

大家好，我是你们的显存优化讲师。今天我们来聊聊重计算技术，也就是常说的激活值重计算。说实话，这个技术是我在训练百亿参数模型时，用得最频繁的「救命稻草」之一。你想想看，当你的显存差那么一点点就爆掉的时候，重计算往往能帮你省下30%-50%的显存。

先说说基本原理。在神经网络前向传播时，每一层都会产生激活值（activation）。这些激活值需要保留到反向传播时使用。比如一个100层的Transformer，每一层都要存下所有中间结果。显存就是这么被吃掉的。

重计算的思路很简单：我不存了，用的时候再算一遍。具体来说，前向传播时只保留一部分关键节点的激活值（比如每隔几层存一个checkpoint）。反向传播时，如果发现某个激活值没存，就从最近的checkpoint重新计算一遍。

核心思想：用时间换空间

重计算本质上是用额外的计算开销，换取显存空间的释放。典型的trade-off是：
显存节省：30%-50%
计算开销增加：20%-33%

我在项目中遇到过这样的情况：一个72层的GPT模型，如果不做重计算，单卡根本放不下。用了重计算后，显存占用从48GB降到了28GB，虽然训练时间多了15%，但至少能跑了。嗯，这就是典型的「有总比没有好」。

不是所有层都值得重计算。我个人习惯把层分成三类：

为什么会这样？Attention层的QKV矩阵维度是[batch, seq_len, hidden_dim]，这个量级非常大。但Attention的计算本身并不重，主要是矩阵乘法。所以重计算Attention的性价比很高。

FFN层就不一样了。FFN的中间维度通常是hidden_dim的4倍，计算量很大。如果你重计算FFN，节省的显存不多，但增加的计算开销却不少。我建议只对FFN的中间激活做重计算，输入输出保留。

我的经验法则：

选择重计算层时，看一个指标：激活值大小 / 计算量。这个比值越大，越值得重计算。我一般设定阈值为0.3，比值超过0.3的层就做重计算。

这里有个坑，我曾经踩过。流水线并行本身就有多个stage，每个stage负责几层。如果你在每个stage内部再做重计算，需要特别注意通信和计算的调度。

举个例子，假设流水线有4个stage，每个stage有6层。你决定每3层存一个checkpoint。那么反向传播时，每个stage需要从checkpoint重新计算中间3层的激活值。问题来了：重计算和流水线气泡怎么共存？

我建议的做法是：

避坑指南：

我曾经在一个8卡流水线并行训练中，把每个stage的checkpoint设成了每层一个。结果重计算导致流水线气泡从20%飙升到了45%，训练效率反而下降了。后来我改成每3层一个checkpoint，气泡才回到25%。

重计算不是免费的午餐。我们来算一笔账：

假设一个Transformer层的前向传播需要1个单位时间，反向传播需要2个单位时间。如果不做重计算，总时间是3个单位。如果做全量重计算（每层都重算），前向传播变成2个单位（一次正常前向，一次重计算），反向传播还是2个单位，总时间变成4个单位。计算开销增加了33%。

但实际中，我们很少做全量重计算。选择性重计算通常只增加20%-25%的计算开销。这个代价换来的显存节省，在大多数场景下是值得的。

我记得有一次训练一个130亿参数的模型，显存刚好差5GB。我用了选择性重计算，只重计算了Attention层和FFN的中间激活，显存降了8GB，计算开销只增加了18%。嗯，这个trade-off我很满意。

计算开销的量化分析：

假设模型有L层，每层前向计算量为F，反向计算量为2F：

最后说一个实用技巧：动态调整重计算策略。我在训练过程中会监控显存使用率和计算时间。如果显存还有余量，就减少重计算层数，提高训练速度。如果显存快爆了，就增加重计算层数。说白了，就是让系统自己找到最优的平衡点。

好了，重计算技术就讲到这里。记住核心：不是所有激活值都值得存，也不是所有层都值得重算。找到那个平衡点，你的训练效率就能最大化。

想要更多资料，联系可可