Memory-Efficient Activation Checkpointing with Sliding Window and Hirschberg's Algorithm for 0/1 Knapsack Solving in PyTorch
Not on the map. Why?
Abstract (arXiv)
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, dpknapsack, allocates a full dynamic programming (DP) table of shape (n+1) × (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 dpknapsackslidinghirschberg, 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 dpknapsack fails at n = 100, a 20× increase in computable problem size. In addition, our benchmarks show a consistent 25-28% runtime speedup over dpknapsack. The implementation is merged into PyTorch and released in version 2.10.
Key facts. No citing papers in the map yet (data through 14 Oct 2026). Builds on 17 of its 35 references that are in the map. First posted on arXiv 9 Aug 2026.
Builds on
17 of 35 references.
References · 17
- 2019PyTorch: An Imperative Style, High-Performance Deep Learning Librarycited by 20,305
- 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,673
- 2019Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelismcited by 1,868
- 2020Reformer: The Efficient Transformercited by 1,668
- 2019ZeRO: Memory Optimizations Toward Training Trillion Parameter Modelscited by 1,301
- 2016Training Deep Nets with Sublinear Memory Costcited by 848
- 2021ZeRO-Offload: Democratizing Billion-Scale Model Trainingcited by 341
- 2022Reducing Activation Recomputation in Large Transformer Modelscited by 302
- 2017The Reversible Residual Network: Backpropagation Without Storing Activationscited by 217
- 2019Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerializationcited by 91
- 2016Memory-Efficient Backpropagation Through Timecited by 78
- 2021Torch.fx: Practical Program Capture and Transformation for Deep Learning in Pythoncited by 40
5 more
- 2019Optimal checkpointing for heterogeneous chains: how to train deep neural networks with limited memorycited by 22
- 2019A Graph Theoretic Framework of Recomputation Algorithms for Memory-Efficient Backpropagationcited by 11
- 2025A Survey on Memory-Efficient Transformer-Based Model Training in AI for Sciencecited by 3
- 2022XEngine: Optimal Tensor Rematerialization for Neural Networks in Heterogeneous Environmentscited by 2
- 2022Speeding Hirschberg Algorithm for Sequence Alignmentcited by 1
Built on by
Not cited yet.
Similar, not cited
Close in wording, no citation either way
- 2022Few-Bit Backward: Quantized Gradients of Activation Functions for Memory Footprint Reduction
- 2025An Efficient Compression of Deep Neural Network Checkpoints Based on Prediction and Context Modeling
- 2026FlashOptim: Optimizers for Memory-Efficient Training
- 2020Check-N-Run: A Checkpointing System for Training Deep Learning Recommendation Models ↗
- 2021Pro-KD: Progressive Distillation by Following the Footsteps of the Teacher
- 2024Sparse Checkpointing for Fast and Reliable MoE Training
- 2026FLOP-Efficient Training: Early Stopping Based on Test-Time Compute Awareness
- 2021Accelerating SLIDE Deep Learning on Modern CPUs: Vectorization, Quantizations, Memory Optimizations, and More