Memory-Efficient Activation Checkpointing with Sliding Window and Hirschberg's Algorithm for 0/1 Knapsack Solving in PyTorch
Jędrzej Maczan
cs.LG
Submitted: 2026-08-09
Updated: 2026-08-11
Comments: Accepted to COLM 2026 Workshop on Efficient Reasoning
Code: https://github.com/pytorch/pytorch
License: http://creativecommons.org/licenses/by/4.0/
The gist: Activation checkpointing minimizes the runtime of neural networks under a given memory budget, by selecting which intermediate tensors to store and which to recompute.
Terminology
Abstract
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.
Sources
- Training Deep Nets with Sublinear Memory Cost
- The Reversible Residual Network: Backpropagation Without Storing Activations
- Memory-Efficient Backpropagation Through Time
- Optimal checkpointing for heterogeneous chains: how to train deep neural networks with limited memory
- Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization
- Dynamic Tensor Rematerialization
- Reformer: The Efficient Transformer
- Reducing Activation Recomputation in Large Transformer Models
- A Graph Theoretic Framework of Recomputation Algorithms for Memory-Efficient Backpropagation
- Speeding Hirschberg Algorithm for Sequence Alignment
- Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM
- PyTorch: An Imperative Style, High-Performance Deep Learning Library
- Torch.fx: Practical Program Capture and Transformation for Deep Learning in Python
- ZeRO-Offload: Democratizing Billion-Scale Model Training
- Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks