Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization

Paras Jain; Ajay Jain; Aniruddha Nrusimha; Amir Gholami; Pieter Abbeel; Kurt Keutzer; Ion Stoica; Joseph E. Gonzalez

Source-linked AI summary

Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization

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

arXiv:1910.02653v3Oct 7, 2019cs.LGcs.CVcs.DCstat.ML

TL;DR

DNN training memory is dominated by activation tensors, while prior rematerialization methods rely on restrictive assumptions for graph structure and cost. Checkmate formulates tensor rematerialization as a hardware-aware MILP and provides a two-phase LP-rounding approximation. It achieves minimal computational overhead across memory budgets and enables larger inputs, including up to 5.1× higher batch sizes on the same accelerator.

Problem
Activation tensors dominate DNN training memory, while prior rematerialization approaches are limited by linear-graph, uniform-cost, and non-rematerialized-gradient assumptions.
Method
Checkmate formulates rematerialization as a MILP with a flexible search space and provides a two-phase deterministic LP-rounding approximation, using hardware-aware profile-guided costs.
Results
5.1× higher batch sizes are enabled on the same accelerator, while optimal rematerialization has minimal computational overhead across a wide range of memory budgets.
Takeaways & Limitations
Checkmate enables larger models, larger batch sizes, and more complex signals with minimal computation overhead.
Takeaways & Limitations
ILP solving is NP-hard and may be infeasible for neural networks with hundreds of layers, motivating the approximation algorithm.

Abstract

from arXiv · show

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.

1 INTRODUCTION

DNN training is constrained by activation-memory demand, while prior rematerialization methods are limited to simpler graphs and uniform costs. Checkmate formulates rematerialization as an optimization problem and provides solver-based and approximation methods for reducing memory with minimal overhead.

Activation tensors needed for backpropagation dominate training memory in workloads using high-resolution images, 3D point clouds, and long language sequences.
The memory wall created by limited high-bandwidth device memory restricts exploration of novel DNN architectures.
Discarding selected activations and rematerializing them when needed lets large DNNs fit within memory at the cost of additional computation.
21GB of memory is saved for the illustrated 32-layer network by freeing activations early, with a modest runtime increase from rematerialization.
Prior heuristics assume linear graphs, equal node costs, and no gradient rematerialization, limiting their efficiency and generality for nonlinear DNNs.
Checkmate formalizes rematerialization as a mixed integer linear program and adds a two-phase deterministic LP-rounding approximation algorithm.
5.1× larger input sizes are enabled at minimal overhead, while experiments report substantially reduced memory usage and minimal computational overhead across image classification and semantic segmentation architectures.

2 MOTIVATION

Training memory is dominated by intermediate activations, creating a memory wall that limits DNN design and motivates rematerialization. Checkmate targets arbitrary graphs and variable layer costs while minimizing recomputation within a memory budget.

Memory limits: Training memory is dominated by intermediate activation tensors whose size depends on input dimensions, rather than parameters and gradients.This is especially important for high-resolution images, 3D point clouds, and long natural-language sequences.
Rematerialization: Rematerialization discards selected forward-pass tensors and recomputes them during backpropagation, trading additional computation for lower memory usage.The goal is to fit arbitrary networks within a memory budget while incurring minimal runtime overhead.
Limits of prior work: Prior checkpointing heuristics assume linear graphs, equal node costs, and no gradient rematerialization, limiting their efficiency and applicability to nonlinear DNNs.Residual connections and architectures such as U-Net cannot generally be handled efficiently under these assumptions.
Checkmate formulation: Checkmate searches schedules for arbitrary graphs with variable memory and computation costs, subject to dependency correctness and a RAM budget.Its objective minimizes additional computation cost using accelerator-specific or profiled layer costs when available.

3 RELATED WORK

Prior approaches address memory through checkpointing, reversible networks, distributed computation, or activation compression, but each has scope or trade-offs. Checkmate instead rematerializes entire data-flow subgraphs while accounting for realistic graph and layer costs.

Checkpointing: Checkpointing heuristics provide memory reductions for idealized linear, unit-cost graphs but are not optimal when DNN layers vary in memory and computation cost.Earlier methods include O(√n) and O(log n) memory strategies, with optimality proved only for restricted linear graphs.
Checkpointing: Prior checkpointing methods are limited by linear-graph assumptions, whereas Checkmate handles nonlinear data-flow graphs and can recompute entire subgraphs.Register rematerialization is generally restricted to cheaply recomputable values with register-resident dependencies.
Reversible networks: Reversible networks reduce memory by recomputing intermediate values from later forward results, and are described as complementary to rematerialization.The paper reports that rematerialization enables greater savings and a wider range of memory budgets than reversibility in its comparison.
Distributed computation: Distributed computation and gradient accumulation can address memory constraints, but may require additional accelerators, fast networks, partitioning, or multiple sub-batches.Gradient accumulation can also degrade performance when batch-normalization statistics are affected.
Activation compression: Activation compression reduces memory by compressing or quantizing activations, but can reduce accuracy; rematerialization is described as incurring no accuracy penalty.The paper characterizes its approach as mathematically equivalent to rematerialization-free training in this respect.

4 OPTIMAL REMATERIALIZATION

Checkmate formulates memory-constrained tensor rematerialization for general data-flow graphs as a MILP that jointly schedules computation, retention, recomputation, and garbage collection. The formulation enforces dependency correctness and memory feasibility, while frontier-advancing stages improve tractability.

Problem formulation: Checkmate minimizes execution time for general data-flow graphs while guaranteeing that schedules remain within device memory limits.The rematerialization problem is formulated as a MILP solvable with commercial or open-source standard solvers.
Problem formulation: Each graph node has a memory cost M_v and computation cost C_v, and the schedule must compute the terminal node under M_budget at minimum total computation cost.Nodes represent tensor-producing operations and edges represent dependencies.
Schedule representation: The schedule is unrolled into T stages, with binary variables indicating retained values and recomputed operations; setting T = n yields O(n^2) operations and constant memory for linear graphs.This representation generalizes checkpointing by allowing values to be retained and deallocated multiple times.
Memory constraints: Memory accounting variables U_t,k track memory after evaluating node v_k, while FREE variables model dependency deallocation subject to future-use and checkpoint constraints.A dependency may be freed after evaluation only when no relevant future use requires it.
Memory constraints: The MILP linearizes polynomial garbage-collection conditions using binary-polynomial and indicator-constraint reformulations.The resulting constraints prevent double deallocation and preserve feasibility across stages.
Tractability via frontier-advancing stages: Frontier-advancing stages reduce the feasible set and can dramatically improve optimization time without changing the objective.For an 8-layer linear graph, Gurobi improved from 9.4 hours to 0.23 seconds with the same objective.

5 APPROXIMATION

The approximation method relaxes the rematerialization MILP to an LP, then uses deterministic two-phase rounding to obtain feasible near-optimal schedules in polynomial time. Budget feasibility requires reserving an allowance because rounding preserves correctness constraints but may increase memory.

Approximation motivation: ILP solving is practical for many instances, with solve times from seconds to an hour and less than 1% overhead relative to long training runs.However, ILP solving is NP-hard and can become impractical for very deep architectures.
5.1 Relaxing integrality constraints: LP relaxation makes the rematerialization problem polynomial-time solvable and provides a lower bound on the optimal integral cost.The relaxation replaces binary decision variables with continuous values in [0, 1].
5.2 Two-phase rounding: Direct deterministic rounding can violate constraints, while randomized rounding failed to find a feasible VGG16 schedule in 50,000 samples.The randomized experiment used an LP relaxation at a budget four times smaller than default.
5.2 Two-phase rounding: Two-phase rounding deterministically rounds the checkpoint matrix, then conditionally solves the binary computation matrix to correct violated correctness constraints.The algorithm starts with an all-zero computation matrix and adds only computations required to satisfy constraints, preserving already satisfied constraints.
Budget feasibility: Rounding preserves correctness but may violate the memory budget, so the method reserves an allowance using Ut,k ≤ (1 − ϵ)Mbudget, with ϵ = 0.1 working well empirically.The allowance compensates for memory increases introduced by rounding.

6 EVALUATION

The evaluation compares Checkmate with checkpointing baselines across image classification and semantic segmentation networks, measuring memory–overhead trade-offs, larger feasible batches, and approximation quality. Checkmate produces faster in-budget schedules, supports substantially larger batches, and its rounding method remains close to optimal.

Evaluation setup: The evaluation tests rematerialization across VGG16, VGG19, ResNet50, MobileNet, U-Net, FCN with VGG layers, and SegNet.The experiments examine memory–overhead trade-offs, large inputs, and approximation to optimal policies.
6.3 What is the trade-off between memory usage and computational overhead?: At the 16 GB V100 budget, Checkmate achieves a 1.20× speedup over linearized greedy and a 1.38× speedup over linearized √n on U-Net.Figure 5 measures computational overhead relative to the best unrestricted strategy using profiled costs on one NVIDIA V100.
Baselines: Checkmate’s feasible set is a superset of baseline heuristics, while baselines fix checkpoint decisions and optimize their recomputation schedules.This comparison uses the same lowest-cost recomputation procedure after each baseline policy is specified.
6.3 What is the trade-off between memory usage and computational overhead?: Checkmate enables U-Net batch size 32 with less than 10% higher overhead, compared with 23 GB required without rematerialization or the original baselines.Across the evaluated networks, it produces faster execution plans than the compared checkpointing methods.
6.4 Are large inputs practical with rematerialization?: Checkmate increases MobileNet batch size to 1105, or 1.73× over the best greedy baseline and 5.1× over caching all activations.The same schedules can increase image resolution instead of batch size.

7 CONCLUSION

The paper frames limited accelerator memory as a central constraint on training large neural networks and presents Checkmate as a rematerialization-based response.

Memory capacity, rather than data and compute availability, is identified as the current bottleneck for state-of-the-art model development.
Checkmate supports general nonlinear computation graphs, non-uniform tensor costs, and hardware-aware profile-guided optimization.
Checkmate enables large models to be trained with limited available memory.
Evaluation reports minimal computational overhead across a wide range of memory budgets and substantially larger batch sizes for high-resolution models.

A INTEGRALITY GAP

The integrality gap measures how much the optimal integer-program objective exceeds its fractional relaxation, helping explain solver performance.

The integrality gap is the maximum ratio between the optimal ILP value and the value of its relaxation.
For a problem instance I, the formulation uses I = (G, C, M, Mbudget).
Because the ILP is a minimization problem, COSTint ≥ COSTfrac and IG ≥ 1.
21.56 to 1.18: frontier-advancement sharply reduces the integrality gap for the 8-layer linear neural network graph.

B.1 AP √n and AP greedy

The AP approach uses articulation points as checkpoint candidates, then computes recomputation plans that preserve dependency correctness under memory constraints.

Articulation points are graph vertices whose removal increases the number of connected components and can be found in O(V + E).
Articulation points are checkpoint candidates because later topological vertices have no dependencies on earlier vertices before the articulation point.
The method applies Chen’s heuristics to select checkpoints and solves for an optimal recomputation plan R to restore correctness.
With checkpoints fixed, R can be obtained in O(|V||E|) using a graph traversal for each row of R.

B.2 Linearized √n and Linearized greedy

The linearized baselines replace the forward graph with edges between consecutive topologically ordered vertices, then optimize recomputation using the resulting checkpoint set.

Linearized √n and Linearized greedy: The forward graph Gfwd is converted into a linear graph Glin by connecting consecutive vertices in topological order.
Linearized √n and Linearized greedy: The comparison evaluates memory usage and computational cost in milliseconds under a profile-based cost model.
Linearized √n and Linearized greedy: The figure compares two-phase LP rounding with randomized and deterministic rounding of S∗ across different models.
Linearized √n and Linearized greedy: Glin does not properly encode data dependencies, but it provides a linear graph that baseline methods can analyze.
Linearized √n and Linearized greedy: After a baseline produces checkpoints, the method forms checkpoint matrix S and finds the optimal recomputation plan R.

C HARDNESS OF REMATERIALIZATION

Rematerialization decision problems closely relate to register allocation on computation DAGs, while constrained optimization formulations provide a practical way to decide feasibility and can remain tractable for many networks.

Sethi (1973) reduced 3-SAT to a register-allocation decision problem for straight-line programs without recomputation.The associated computation graphs are result-rooted DAGs whose nodes represent operations and whose edges carry values.
Register allocation can be represented as stone or register moves on DAG nodes, including computing leaves, removing registers, and computing a node from its children.A node may also be computed by moving a stone from one of its children when the required child registers are available.
RP-DEC asks whether a result-terminated data-flow DAG has an execution plan using at most memory b and cost at most |V| under unit computation and storage costs.The problem specifies a leaf terminal node t, a memory bound b, and unit cost and memory for each node result.
RP-DEC is decidable by solving a sufficiently staged memory-constrained optimization problem and checking whether the resulting plan costs at most |V|.Its formulation differs subtly from register allocation because data-flow graphs terminate at the result and register computations may operate in place.
A 0-1 integer program for optimal allocation under an instruction schedule had empirical complexity O(n2.5), while the frontier-advancing constrained optimization problem is tractable for many networks.These observations qualify the theoretical hardness with evidence that related optimization formulations can be practically manageable.

D COMPARISON OF APPROXIMATIONS

The comparison evaluates randomized and deterministic two-phase rounding of the LP-relaxation matrix S∗, finding that deterministic rounding consistently yields lower-cost schedules.

Two-phase deterministic rounding of S∗ produces consistently lower-cost schedules than two-phase randomized rounding.Randomized rounding instead produces a range of feasible solutions.