先这样答
核心取舍是用计算时间换取显存空间。激活检查点机制改变了标准的反向传播流程,不再于前向传播时保存所有层的中间激活值,而是仅保存部分关键节点的激活状态作为检查点。当反向传播计算到未保存激活值的层时,系统会从最近的检查点出发,重新执行一次前向计算来获取这些数据。
这种机制带来了直接的显存收益。通过分段重算,激活值带来的显存占用可以从 O(层×序列长) 的线性增长规模,降低到 O(√层) 量级。在代价方面,由于反向传播过程中需要重算缺失的激活值,这相当于多进行了一轮局部的前向计算。在实际训练中,这种重复计算通常会导致整体吞吐量下降两三成。
尽管存在吞吐量下降的代价,但省出的显存空间换取了训练更大 batch size 或更长序列的可行性。特别是在长上下文模型训练中,由于序列长度增加导致激活显存激增,激活检查点已经成为保证训练能跑通的必开选项。
在实际工程落地中,我们通常会采用选择性重算策略,即只重算那些计算代价便宜但占用显存大的计算段。同时,这项技术也会与 ZeRO 显存优化、FlashAttention 等机制叠加使用。工程调优的核心,就是通过测试找到重算哪些计算段的最优集合,从而在计算吞吐和显存限制之间达到最佳平衡。
面试官会怎么追问
-
「既然吞吐会下降,那什么场景下必须开激活检查点?」 当模型参数量、batch size 或序列长度的设定导致显存溢出时必须开启。特别是在长上下文训练场景下,激活值占用的显存随序列长度线性增长,此时开启激活检查点是保证训练能够正常运行的基础条件。
-
「你提到的选择性重算是怎么做的?」 标准重算是按网络层进行切分,而选择性重算是对层内的具体算子进行区分。做法是优先重算那些计算量小但激活值显存占用大的操作,保留那些计算成本高但激活值小的部分,以此在时间代价和显存收益间寻找更优的折中。
-
「激活检查点和 FlashAttention 配合使用时有什么关系?」 两者是互补关系。FlashAttention 通过分块计算避免了注意力矩阵的实例化,已经大幅降低了注意力机制本身的激活显存。两者叠加使用时,激活检查点可以进一步压缩网络其他部分前向过程留存的激活值,两者配合是当前大模型训练的标准配置。
回答的坑
- 认为激活检查点可以减少整体训练时间。正确方向是明确指出它是用时间换空间,必然会增加单步训练耗时,其核心目的是为了跑通原本会超出显存限制的训练配置。
- 把激活检查点和模型参数检查点混淆。正确方向是强调激活检查点是训练过程中的显存优化机制,针对的是前向传播产生的中间激活值,而不是用于容错恢复的模型权重保存。
同系列的题