浅谈ai - Activation Checkpointing - 时间换空间

前言

曾在游戏世界挥洒创意，也曾在前端和后端的浪潮间穿梭，如今，而立的我仰望AI的璀璨星空，心潮澎湃，步履不停！愿你我皆乘风破浪，逐梦星辰！

Activation Checkpointing（激活检查点），在 DeepSpeed 里是一个非常实用的显存优化技术。简单来说，它的核心目标就是：

“节省显存，代价是多做点计算。”

下面我给你讲明白它是啥、怎么干的、为什么有效、啥时候该用。

---

🧠 一句话解释

Activation Checkpointing 就是：

在前向传播时不保存所有中间激活（activation），而是只保存“关键点”的激活；
等到反向传播需要时，再重新计算中间丢掉的部分。

---

📦 举个例子（以 Transformer 为例）

假设你有一个 12 层的 Transformer：

• 正常训练时，前向传播完 12 层，会保存 每一层的激活输出（用于反向传播时算梯度）。

• 这些激活可能占你一大半的显存！

💡 但是，如果你设置了 activation checkpointing：

• 你只保存第 0、4、8、12 层的激活（比如每隔 4 层存一个“检查点”）。

• 当反向传播需要第 2 层激活时，Deepspeed 会：

  • 重新从第 0 层开始计算到第 2 层。

  • 得到激活，然后继续正常反传。

• 这样就省下了很多