浅谈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 层。
-
得到激活，然后继续正常反传。
-
-
这样就省下了很多显存，但代价是计算时间会稍微变长（因为要“回头重算”）。
🔍 技术实现机制
| 步骤 | 做了什么 |
|---|---|
| 前向传播 | 只保存“检查点”处的激活，其他层的激活丢掉 |
| 反向传播 | 当需要中间层激活时，重新从上一个检查点 forward 一遍 |
| PyTorch 实现 | 通常使用 torch.utils.checkpoint.checkpoint(function, *inputs) 来做 |
| Deepspeed 支持 | 可以自动为模型的某些模块启用 checkpoint，比如 TransformerBlock |
🧮 显存节省 vs 计算开销
| 优点 | 缺点 |
|---|---|
| 大幅减少激活存储显存（最多可省 50%+） | 增加前向传播的计算量（因为反向时要重算一部分） |
| 可以训练更大的模型 / 更大 batch | 训练速度略慢（但通常接受） |
典型使用场景：
-
显存紧张，训练不了大的 batch 或模型。
-
训练 GPT/BERT 这类“块状重复”的模型时特别有效（容易切分 block）。
✅ 在 Deepspeed 中如何启用？
在 deepspeed_config.json 中这样配：
{
"activation_checkpointing": {
"partition_activations": true,
"contiguous_memory_optimization": true,
"cpu_checkpointing": false,
"number_checkpoints": 4
}
}
| 配置项 | 含义 |
|---|---|
partition_activations | 按 rank 切分激活，进一步省显存 |
contiguous_memory_optimization | 内存优化，避免碎片化 |
cpu_checkpointing | 是否把 checkpoint 存 CPU（更慢，慎用） |
number_checkpoints | 控制多少层之间插一个 checkpoint（类似“每几层存一次”） |
📌 总结一句话
Activation Checkpointing 是一种“用时间换空间”的策略，省显存的同时代价是略微增加计算。对于大模型（如 GPT、BERT）训练来说是非常常见的标配技术。
如果你要训练 13B、30B 这种大模型或者 batch 太大撑不住，那这技术几乎是必开项了。