基于重放契约与选择性保留的状态感知激活重计算

激活重计算通过在反向传播期间重放前向计算片段来节省加速器内存，但只有当重放所见的状态与原始前向传播所见的状态一致时，重放才是正确的。可变缓冲区、存在别名的字节以及随机数生成器状态可能在此期间发生变化，因此，重新执行未修改的代码可能会对不同的计算求导。我们将这一状态显式化。带版本的重放契约规定重放必须恢复哪些内容，联合规划器则针对每个片段，在保留激活、保存整个缓冲区的快照和选择性保留区域之间进行选择，并按实测的内存开销与恢复成本评估每个选项。独立的导数检查和字节级检查用于验证该契约，一组 Torch/JAX 实验保留了全部 5,900 条执行记录，包括匹配、拒绝和不一致的情况。在配置固定的稀疏写入 H100 工作负载上，每项测量均使用五十个全新进程；在 44 GiB 的预算下，相较于成本核算正确的即时／整缓冲区规划器，联合规划使配对测量的单步耗时降低了 18.16%（近似 97.5% 比值区间为 $[.8172,.8195]$）。在 32 GiB 的预算下，即时／整缓冲区规划器不存在可行的静态策略，而联合规划比固定区域策略快 4.98%，内存用量为 29.94 GiB，而非 26.20 GiB。一个完整的 CIFAR-100 密集队列应用在六个训练随机种子下运行了七种方法，每种方法均进行 100 轮编码器训练和 25 轮探测器训练，同一种子内各方法的测试得分在字节级完全一致。在该应用中，联合规划的已分配内存比整缓冲区快照多 22.41%，进程耗时的几何平均比值为 1.0080（各随机种子的取值范围为 $[.9845,1.0191]$）。这些结果共同揭示了选择性保留何时能够带来收益：具有大型激活的稀疏写入场景可以受益，而在密集队列工作负载中，整缓冲区快照使用的已分配内存更少。