优化PyTorch的激活检查点，让大模型训练更省内存、更快。
Memory-Efficient Activation Checkpointing with Sliding Window and Hirschberg's Algorithm for 0/1 Knapsack Solving in PyTorch
- 用滑动窗口+希尔斯伯格算法替代传统DP表，降低内存占用。
- 在64GB内存下可处理2000个操作，是原方法的20倍。
- 提速25%-28%，已集成到PyTorch 2.10中。
激活检查点在给定内存预算下通过选择存储哪些中间张量、哪些重算来最小化神经网络运行时间。PyTorch将此问题建模为0/1背包问题，其中前向-反向计算图中的操作作为物品，具有内存开销（重量）和运行时间节省（价值）。默认求解器dp_knapsack需分配大小为$(n+1) imes (W+1)$的完整动态规划表，其中$n$为操作数，$W$为量化内存预算。该方法资源消耗大，在64GB RAM机器上$ n = 100 $时即崩溃。本文提出dp_knapsack_sliding_hirschberg，结合滑动窗口技巧与希尔斯伯格算法，将峰值内存从$O(nW)$降至$O(W)$，同时保持精确最优解。实验表明，可在$ n = 2000 $时成功求解，而dp_knapsack在$ n = 100 $时已失败，计算规模提升20倍；基准测试显示相较dp_knapsack持续获得25%-28%的运行时间加速。该实现已合并至PyTorch并随2.10版本发布。
原文摘要 · Abstract (English)
Activation checkpointing minimizes the runtime of neural networks under a given memory budget, by selecting which intermediate tensors to store and which to recompute. PyTorch solves this as a 0/1 knapsack problem, where operations from a joint forward-backward computation graph are items with a memory cost (weight) and a runtime saving (value). The default solver, dp_knapsack, allocates a full dynamic programming (DP) table of shape $(n+1) \times (W+1)$, where $n$ is the number of operations and $W$ is the quantized memory budget. This method is resource-hungry and crashes at $n = 100$ items on a machine with 64 GB RAM. In this paper, we introduce dp_knapsack_sliding_hirschberg, which combines the sliding window trick and Hirschberg's algorithm to reduce peak memory from $O(nW)$ to $O(W)$ while preserving the exact optimal solution. Our experiments show successful knapsack execution at $n = 2000$, where dp_knapsack fails at $n = 100$, a 20$\times$ increase in computable problem size. In addition, our benchmarks show a consistent 25-28\% runtime speedup over dp_knapsack. The implementation is merged into PyTorch and released in version 2.10.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。