# Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization

arXiv: https://arxiv.org/abs/1910.02653 · HTML: https://mlsanity.com/paper/1910.02653/

## Key facts

- Cited by 91 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 7 Oct 2019.

## Built on by (most cited first, 30 of 91)

- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://mlsanity.com/paper/2104.04473/) (2021)
- [Putting NeRF on a Diet: Semantically Consistent Few-Shot View Synthesis](https://mlsanity.com/paper/2104.00677/) (2021)
- [ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning](https://mlsanity.com/paper/2104.07857/) (2021)
- [Combined Scaling for Zero-shot Transfer Learning](https://mlsanity.com/paper/2111.10050/) (2021)
- [AI and Memory Wall](https://mlsanity.com/paper/2403.14123/) (2024)
- [Software-Hardware Co-design for Fast and Scalable Training of Deep Learning Recommendation Models](https://mlsanity.com/paper/2104.05158/) (2021)
- [TeraPipe: Token-Level Pipeline Parallelism for Training Large-Scale Language Models](https://mlsanity.com/paper/2102.07988/) (2021)
- [ActNN: Reducing Training Memory Footprint via 2-Bit Activation Compressed Training](https://mlsanity.com/paper/2104.14129/) (2021)
- [Stateful Large Language Model Serving with Pensieve](https://mlsanity.com/paper/2312.05516/) (2023)
- [Dorylus: Affordable, Scalable, and Accurate GNN Training with Distributed CPU Servers and Serverless Threads](https://mlsanity.com/paper/2105.11118/) (2021)
- [Bamboo: Making Preemptible Instances Resilient for Affordable Training of Large DNNs](https://mlsanity.com/paper/2204.12013/) (2022)
- [Differentially Private Bias-Term Fine-tuning of Foundation Models](https://mlsanity.com/paper/2210.00036/) (2022)
- [Compute-Efficient Deep Learning: Algorithmic Trends and Opportunities](https://mlsanity.com/paper/2210.06640/) (2022)
- [GACT: Activation Compressed Training for Generic Network Architectures](https://mlsanity.com/paper/2206.11357/) (2022)
- [Accelerating Quadratic Optimization with Reinforcement Learning](https://mlsanity.com/paper/2107.10847/) (2021)
- [DISTFLASHATTN: Distributed Memory-efficient Attention for Long-context LLMs Training](https://mlsanity.com/paper/2310.03294/) (2023)
- [Hidet: Task-Mapping Programming Paradigm for Deep Learning Tensor Programs](https://mlsanity.com/paper/2210.09603/) (2022)
- [POET: Training Neural Networks on Tiny Devices with Integrated Rematerialization and Paging](https://mlsanity.com/paper/2207.07697/) (2022)
- [KAISA: An Adaptive Second-Order Optimizer Framework for Deep Neural Networks](https://mlsanity.com/paper/2107.01739/) (2021)
- [Alpa: Automating Inter- and Intra-Operator Parallelism for Distributed Deep Learning](https://mlsanity.com/paper/2201.12023/) (2022)
- [On-device Training: A First Overview on Existing Systems](https://mlsanity.com/paper/2212.00824/) (2022)
- [Sequential Aggregation and Rematerialization: Distributed Full-batch Training of Graph Neural Networks on Large Graphs](https://mlsanity.com/paper/2111.06483/) (2021)
- [TinyTrain: Resource-Aware Task-Adaptive Sparse Training of DNNs at the Data-Scarce Edge](https://mlsanity.com/paper/2307.09988/) (2023)
- [The CoRa Tensor Compiler: Compilation for Ragged Tensors with Minimal Padding](https://mlsanity.com/paper/2110.10221/) (2021)
- [Unicron: Economizing Self-Healing LLM Training at Scale](https://mlsanity.com/paper/2401.00134/) (2023)
- [PatrickStar: Parallel Training of Pre-trained Models via Chunk-based Memory Management](https://mlsanity.com/paper/2108.05818/) (2021)
- [A Comprehensive Survey on Distributed Training of Graph Neural Networks](https://mlsanity.com/paper/2211.05368/) (2022)
- [Asteroid: Resource-Efficient Hybrid Pipeline Parallelism for Collaborative DNN Training on Heterogeneous Edge Devices](https://mlsanity.com/paper/2408.08015/) (2024)
- [Dissecting the Runtime Performance of the Training, Fine-tuning, and Inference of Large Language Models](https://mlsanity.com/paper/2311.03687/) (2023)

## Abstract (arXiv)

We formalize the problem of trading-off DNN training time and memory requirements as the tensor rematerialization optimization problem, a generalization of prior checkpointing strategies. We introduce Checkmate, a system that solves for optimal rematerialization schedules in reasonable times (under an hour) using off-the-shelf MILP solvers or near-optimal schedules with an approximation algorithm, then uses these schedules to accelerate millions of training iterations. Our method scales to complex, realistic architectures and is hardware-aware through the use of accelerator-specific, profile-based cost models. In addition to reducing training cost, Checkmate enables real-world networks to be trained with up to 5.1x larger input sizes. Checkmate is an open-source project, available at https://github.com/parasj/checkmate.