梯度检查点

定义

梯度检查点(Gradient Checkpointing,又称激活重计算)是一种以计算换显存的训练技术:前向传播时只保留少数检查点处的激活,其余中间结果在反向传播需要时从最近的检查点重新前向计算得到。

原理

训练显存的大头往往不是参数而是激活,其随层数与序列长度线性甚至平方增长。对 n 层网络每隔根号 n 层设一个检查点,可把激活显存从 O(n) 降到 O(根号n),额外代价约为一次前向计算(整体训练慢约 20% 到 35%)。

应用

几乎所有大模型训练框架(PyTorch、DeepSpeed、Megatron)都内置该功能;Megatron 的选择性重计算只对注意力等「显存大、计算便宜」的部分重算,进一步优化取舍。在长上下文训练和消费级显卡微调场景中,梯度检查点常与 LoRA、ZeRO 组合使用。