Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization

By Paras Jain, Ajay Jain, Aniruddha Nrusimha, A. Gholami, Pieter Abbeel, Kurt Keutzer, Ion Stoica, Joseph E. Gonzalez. 2019

244 citations recorded by the source corpus.

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 this https URL.

Subfields: Artificial Intelligence, Computational Mathematics, Hardware and Architecture, Efficiency & Compression, Optimization & Training

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computational Mathematics

In Hardware and Architecture

In Efficiency & Compression

In Optimization & Training

Citations in the corpus

Train Large, Then Compress: Rethinking Model Size for Efficient Training and Inference of Transformers (2020, background): “Instead, memory storage/movement is the limiting factor in image classi-ﬁcation (Gomez et al., 2017), semantic segmentation (Chen et al., 2017), language modeling (Kitaev et al., 2020), and other tasks (Jain et al., 2020).”
ZeRO: Memory optimizations Toward Training Trillion Parameter Models (2019, background): “1) Reducing Activation Memory: Multiple efforts have focused on reducing the memory footprint of activations through compression [15], activation checkpointing [8], [16], or live analysis [17].”
ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep learning (2021, background): “Multiple efforts have focused on reducing the memory required by activations through compression [28], activation checkpointing [29, 30], or live analysis [31].”
Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM (2021, background): “Activation recomputation [12, 18, 20, 21] is an optional technique that trades off an increase in the number of compute operations performed for additional memory footprint, by running the forward pass a second time just before the backward pass (and stashing only the input activations for a given p”
Memory-Efficient Pipeline-Parallel DNN Training (2020, background): “Activation recomputation [5,9] performs the forward pass a second time during the backward pass to reduce the memory footprint of activation stashes; systems such as Checkmate [14] systematically determine when this should be performed.”
Efficient Memory Management for Large Language Model Serving with PagedAttention (2023, methodology): “Swapping [23, 42, 55], recomputation [7, 24] and their combination [40] have been utilized to reduce the peak memory of training.”
Combined Scaling for Zero-shot Transfer Learning (2021, methodology): “Both our pipelining approach and our SPMD approach make use of gradient checkpointing (Chen et al., 2016), which is also called rematerialization in certain literature (Kumar et al., 2019; Jain et al., 2020).”
Combined Scaling for Open-Vocabulary Image Classiﬁcation (2022, methodology): “Both our pipelining approach and our SPMD approach make use of gradient checkpointing (Chen et al., 2016), which is also called rematerialization in certain literature (Kumar et al., 2019; Jain et al., 2020).”
Software-hardware co-design for fast and scalable training of deep learning recommendation models (2021, background): “DeepSpeed [50] fully shardsmodel parameters, gradients and optimizer states across all nodes, and reconstructs necessary states on the fly using checkpoint partitioning and rematerialization [21, 28] to drastically reduce memory usage.”
Dynamic Tensor Rematerialization (2020, methodology): “DTR has close ties to garbage collection techniques from the PL literature [42], and modern accelerators being developed in the Arch community are often memory-constrained [9, 43].”
Putting NeRF on a Diet: Semantically Consistent Few-Shot View Synthesis (2021, methodology): “We delete intermediate MLP activations during rendering and rematerialize them during the backward pass [6, 19].”
Zico: Efficient GPU Memory Sharing for Concurrent DNN Training (2021, background): “Checkmate [19] formulates tensor recomputation into an optimization problem and provides an approximation algorithm to recompute tensors timely.”
High-performance, Distributed Training of Large-scale Deep Learning Recommendation Models (2021, unknown): “Deep-Speed [44] fully shards model parameters, gradients and optimizer states across all nodes, and reconstructs necessary states on the fly using checkpoint partitioning and rematerialization [17, 24] to drastically reduce memory usage.”
A Survey on Federated Learning for Resource-Constrained IoT Devices (2021, unknown): “Training on edge devices can bene-ﬁt by adapting techniques that have served well in managing this problem in the server context, such as efﬁcient gradient checkpointing [108], tensor rematerialization [109], and recompute [110].”
AI and Memory Wall (2024, unknown): “As such, we need to rethink the training, deployment, and design of AI models as well as how we design AI hardware to deal with this increasingly challenging memory wall.”
Parallel Training of Pre-Trained Models via Chunk-Based Dynamic Memory Management (2021, methodology): “a activation rematerization, first proposed in work [7] and widely investigated in work [12, 13, 16] to trade off activationmemory with more computation and is implemented in PyTorch [28].”