梯度检查点原理与实现
训练大模型的显存困境 训练一个70B参数的模型，仅模型参数（FP16）就需要140GB显存。加上梯度（140GB）、优化器状态（Adam需要280GB）和激活值，总显存需求轻松超过600GB——即使8卡A100 80GB也无法满足。 在这些显存消耗中，激活值是最容易被忽视的部分。在标准反向传播中，前向传播的所有中间激活值都需要保存下来供反向传播计算梯度使用。对于70B模型、batch_size=4、seq_len=4096，激活值占用可达100GB以上。 梯度检查点（Gradient Checkpointing）通过"以时间换空间"的策略，大幅减少激活值的显存占用。 标准反向传播的显存分析 在标准反向传播中，每个操作的前向输出都需要被保存： 前向: x → Layer1 → a1 → Layer2 → a2 → Layer3 → a3 → Loss 保存: a1, a2, a3 (以及各层的中间状态) 反向: Loss ← ∂L/∂a3 ← ∂L/∂a2 ← ∂L/∂a1 ← ∂L/∂x 使用: 重新使用a3计算a2的梯度，重新使用a2计算a1的梯度，... 对于N层Transformer，每层的激活值大小为 O(batch × seq_len × hidden_dim)。所有层的激活值总和为 O(N × batch × seq_len × hidden_dim)，与层数线性增长。 梯度检查点的核心思想 选择性重计算 梯度检查点的核心洞察：不需要保存所有层的激活值，只需保存部分"检查点"层的输入，其余层的激活值在反向传播时重新计算。 检查点设置（每2层一个检查点）: 保存: x, a2, a4, a6, ... 前向: x → L1 → a1 → L2 → a2* → L3 → a3 → L4 → a4* → ... (* = 保存的检查点) 反向: 从a6*出发，重新前向计算a5, a4→计算梯度 从a4*出发，重新前向计算a3, a2→计算梯度 ... 显存-计算权衡 假设N层Transformer，每隔k层设置一个检查点： ...