Gradient Checkpointing: Training Deep Nets with Sublinear Memory Cost

Store a few activations and recompute the rest: square-root memory for one extra forward pass.

Backpropagation keeps every layer's output in memory until the backward pass reaches it, so training memory grows in proportion to depth. This 2016 paper cuts an $n$-layer network into about $\sqrt{n}$ segments, keeps only each segment's input, and recomputes a segment's interior when the backward pass needs it.

On a 1,000-layer ResNet, the best plan without recomputation needs an estimated 48 GB of activations and cannot run on a 12 GB GPU. With recomputation the whole training job runs in 6.6 GB, about 30% slower.

Tianqi Chen, Bing Xu, Chiyuan Zhang and Carlos Guestrin (University of Washington, Dato and MIT) posted this paper to arXiv in April 2016. Chen, Xu and Zhang were among the authors of MXNet, and the paper reads like a systems note from that project: a list of memory optimizations, the algorithms that plan them, and measurements on a ResNet and an LSTM. It is an arXiv preprint with no conference or journal version. The technique it brought to deep learning, now called gradient checkpointing (or activation checkpointing, or activation recomputation), ships in PyTorch, JAX, DeepSpeed and Megatron-LM, usually behind one flag.

The page follows the paper's order. First, why training stores $O(n)$ activations at all. Then two ways to cut that: reusing buffers whose contents are dead, which saves a constant factor, and dropping activations and recomputing them, which changes the growth rate to $O(\sqrt{n})$. Then the planner that applies this to networks whose layers differ in size, the recursion that reaches $O(\log n)$, and what the experiments measured.

Why training memory grows with depth

At inference a network can run in nearly constant memory: once layer 2 has read layer 1's output, that output is never needed again and its buffer can hold layer 3's output. Training cannot do this, because the backward pass needs the forward pass's values.

Take one fully connected layer, $y = Wx$. The backward pass arrives with $\partial L/\partial y$, the gradient of the loss with respect to the layer's output, and must produce two things:

The weight gradient needs $x$, the layer's input as it was during the forward pass. The output $y$ does not contain it (a matrix product mixes the inputs together), so $x$ has to be kept from the moment it is produced until the backward pass returns to this layer. That return happens in reverse order: the last layer's input is used first and the first layer's input last. When the forward pass ends, the input of every layer is still waiting in memory. Some layers keep less: a ReLU's derivative only needs the sign pattern of its input, which its output also carries. The paper's Section 3 asks frameworks to declare the smallest set of values each gradient operator depends on, for this reason.

For scale: a single 64-channel, 56×56 feature map at batch size 32, stored as 32-bit floats, is 32 × 64 × 56 × 56 × 4 bytes = 25.7 MB. The paper's best plan without recomputation, on its 1,000-layer ResNet, comes to about 48.6 GB of feature maps (the paper's term for activation tensors), about 49 MB per layer. Memory grows linearly with depth, so the deepest trainable network is roughly the GPU's memory divided by the per-layer cost: a GPU with twice the memory fits about twice the depth.

Watch the moment the forward sweep ends: both meters peak there, and both peaks are proportional to the number of layers.

Reusing buffers: liveness analysis

MXNet, like Theano and TensorFlow, represents training as a computation graph: each operation is a node, each edge carries a tensor from the operation that writes it to the operations that read it, and the backward pass is added as more nodes in the same graph. Assigning memory to the graph's tensors is the problem a compiler solves when it assigns a program's variables to a limited set of registers, and the paper borrows the compiler's tool for it, liveness analysis. A tensor is live from the operation that writes it until the last operation that reads it. Two tensors whose live ranges do not overlap can use the same buffer.

Two optimizations follow from that rule:

• In-place: an operation writes its output over its input. This is safe only when no other pending operation still reads that input.
• Sharing: a buffer whose tensor has no readers left goes back to a free pool, and a later tensor takes it.

Finding the fewest buffers exactly is a graph-coloring problem: build a graph with one node per tensor and an edge between tensors whose lifetimes overlap, then color it with as few colors as possible. The paper notes that building and coloring this graph costs $O(n^2)$ time. Finding the minimum coloring is NP-complete in general, a textbook result the paper does not state. It uses a one-pass heuristic instead, linear in the number of operations. Every tensor gets a counter equal to the number of operations that still need to read it. The planner visits the operations in topological order (each operation after everything it reads). When an operation runs, it takes a buffer from the free pool if one is there and a new one otherwise, or writes in place over an input whose counter is 1. Then it decrements its inputs' counters, and any tensor whose counter reaches 0 returns its buffer to the pool.

Figure 2 steps through the paper's own example. Five tensors are computed after the input A, and the planner ends with three buffers. Step 2 refuses to compute C in place over B, because B's counter is still 2 (F also reads it). Step 5 does write G over E in place, because G is the last reader of E.

Because MXNet runs this planner once, before execution, rather than as a garbage collector at runtime, it knows the exact size of every buffer in advance. This lets the paper report the exact feature-map memory of a plan without running it, the numbers behind its Figure 5. On the ResNet, in-place plus sharing cut feature-map memory by a factor of two to three. The cost still grows linearly with depth: a 1,000-layer network needs about 48.6 GB of feature maps under this plan, and the authors could train only about 200 layers on their GPU.

How gradient checkpointing works

Buffer reuse cannot go further because every layer's input is live until the backward pass reaches it. The paper's Section 4 drops that assumption: an activation can be discarded during the forward pass if it can be recomputed later from an earlier activation that was kept.

The simplest version is Algorithm 1. Cut the chain of $n$ layers into $k$ contiguous segments. In the forward pass, store only each segment's input (a checkpoint) and let everything inside the segment be overwritten. In the backward pass, take the segments last to first. For each one, reload its checkpoint, run the segment's forward pass again while saving its activations this time, backpropagate through it with those activations, and free them before moving to the previous segment.

# Algorithm 1: backprop with data dropping on a chain of layers
v = input
for seg in segments:                  # FORWARD
    temp[seg] = v                     # keep the segment's input
    for layer in seg:
        v = layer.forward(v)          # internals are overwritten, not kept
g = loss_grad(v, label)
for seg in reversed(segments):        # BACKWARD, last segment first
    v = temp[seg]                     # restore the stored input
    local = {}
    for layer in seg:                 # recompute the segment's activations
        local[layer] = v
        v = layer.forward(v)
    for layer in reversed(seg):       # ordinary backprop inside the segment
        g = layer.backward(g, local[layer])
    # local is freed here: one segment's activations live at a time

Figure 3 runs this on a 16-layer chain with $k = 4$. Press Play and watch the memory trace below the boxes. During the forward pass it climbs slowly, one checkpoint per segment. During the backward pass it rises and falls once per segment, as each segment is recomputed and then freed.

The checkpointed trace peaks at 7 activations instead of 16. It reaches 7 while the last segment is being backpropagated: all 4 checkpoints are held, and the segment's other 3 activations have just been recomputed (its first activation is its own checkpoint). In general that is $n/k + k - 1$, which the paper rounds to $n/k + k$.

Two costs follow from the trace. Time: every discarded activation is recomputed exactly once, during its segment's turn, so the extra work is one forward pass over the whole network. The paper notes the backward pass takes nearly twice as long as the forward pass, so a normal training step is about 3 forward-units, and checkpointing makes it 4, about 33% more in this unit model. The paper measured 30% on a Titan X. Correctness: the recomputation runs the same operations on the same stored inputs, so the gradients equal those of plain backprop. The paper states that all its memory optimizations give "equivalent weight gradient" and can always be applied.

The obvious alternative is to move activations to CPU memory during the forward pass and copy them back for the backward pass. The paper lists that swapping, and model parallelism, as separate approaches that can be combined with its own; recomputation needs no transfers over the CPU-GPU bus (PCI-E), which leaves that bandwidth for data or model parallelism.

Why √n segments minimize peak memory

The number of segments $k$ is the one free choice in Algorithm 1. The paper's equation (1) writes the peak as the cost of backpropagating through the largest segment plus the cost of the stored checkpoints:

With equal segments, the first term is $n/k$ layers and the second is $k$ checkpoints. The two pull in opposite directions. With $k = 1$ there is one segment holding all $n$ layers, which is plain backprop. With $k = n$ every layer's input is a checkpoint, and the $n$ checkpoints cost as much as storing everything. In between, setting the derivative of $n/k + k$ to zero gives the minimum:

At $k = \sqrt{n}$ the two terms are equal. Drag the slider in Figure 4 from one end to the other: the falling and rising dashed curves cross under the bottom of the solid U, and both ends of the slider cost the full $n = 64$.

For the paper's 1,000-layer ResNet, $\sqrt{1000} \approx 31.6$: about 32 segments of about 31 layers, and a peak of about 63 layers' worth of activations instead of 1,000, a factor of 16 if every layer were the same size. Figure 5 of the paper gives about 12 (48.6 GB to 4.0 GB of feature maps) on the real network, whose layers are not the same size. The paper calls the result $O(\sqrt{n})$ memory for feature maps, the activation tensors. Parameters, their gradients and the temporary workspace that operations like convolution allocate are not included and do not shrink.

Real networks: mirror counts and a memory budget

Algorithm 1 has two problems in practice, and the paper names both: the user has to cut the graph into segments by hand and write a custom training loop, and the loop does not combine with the buffer sharing of Section 3. Algorithm 2 fixes both by moving the decision into graph construction. The user gives every node a mirror count $m(v)$. A node with $m(v) = 0$ keeps its output. A node with $m(v) = 1$ has its output dropped, and the graph builder adds a duplicate (a mirror) of the node to recompute that output for the gradient operators that need it. Setting every count to 0 gives the ordinary gradient graph; setting 0 at segment boundaries and 1 inside segments gives Algorithm 1; counts above 1 give the recursive scheme of the next section. Algorithm 2 outputs an ordinary graph plus an execution order, so the liveness planner runs on it unchanged.

With mirror counts available, the first plan the paper tries (Section 4.2) needs no segments: keep the results of expensive operations and drop the cheap ones. In a convolution, batch normalization, ReLU block, keep the convolution output and recompute batch norm, ReLU and pooling, which cost little to run again. On the ResNet this "drop bn-relu" plan saves memory with little extra compute, but it is still linear in depth: about 34.5 GB of feature maps at 1,000 layers.

Equation (1) assumed equal segments. Real layers are not equal: in a ResNet, each stage halves the spatial size and doubles the channels, so an early layer's output is several times larger than a late one's. Cutting every $\sqrt{n}$ layers would give segments of very different sizes. Algorithm 3 cuts by size instead. It walks the graph in order, adds each output's size to a running total, and places a checkpoint when the total exceeds a budget $B$:

# Algorithm 3: memory planning with budget B
temp, x, y = 0, 0, 0
for v in topological_order(graph):
    temp += size_of_output(v)
    if v in candidate_splits and temp > B:
        m[v] = 0                      # keep v's output (a checkpoint)
        x += size_of_output(v)        # x: stored checkpoints
        y = max(y, temp)              # y: largest segment so far
        temp = 0
    else:
        m[v] = 1                      # drop it, recompute in backward
# m is the "mirror plan" handed to Algorithm 2

The two outputs are the two terms of equation (1): $x$ estimates the stored checkpoints and $y$ the largest segment. As printed, the algorithm updates $y$ only when it places a split, so the segment after the last split goes uncounted; Figure 5 below counts it. The paper then scores each plan exactly with the static allocator. Choosing $B$ is the remaining problem, and the paper's appendix handles it in two steps. Run the algorithm once with $B = 0$, which makes every candidate a checkpoint, and set $B = \sqrt{xy}$ from that run. Then try six values of $B$ between $B/\sqrt{2}$ and $\sqrt{2}\,B$ and keep the plan with the lowest exact cost. The authors' unoptimized Python took a few seconds to produce the plans for the experiments.

For a uniform chain of $n$ layers of size 1, the $B = 0$ run gives $x = n$ and $y = 1$, so the starting budget is $\sqrt{n}$, the answer of equation (1). Figure 5 runs Algorithm 3 on two 36-layer chains. On the uniform one, $B_0 = 6$ and the search finds peak 12, equal to $2\sqrt{36}$. On the ResNet-like one (four stages of nine layers with output sizes 8, 4, 2 and 1, 135 units in total), $B_0 = \sqrt{135 \cdot 8} \approx 32.9$ and the best grid point, $B \approx 30.7$, places 4 checkpoints: $x = 21$, $y = 32$, an estimated peak of 53 against 135 for storing everything. Cutting the same chain into six equal six-layer segments gives 67, because the first segment holds six of the largest layers.

At $B = 0$ every layer is a checkpoint and the estimate is the total plus the largest layer, slightly worse than storing everything. At the right end there are no checkpoints and the estimate is the total again. On the ResNet-like chain, the best budgets put splits close together in the large early stage and far apart in the small late stages, which equal-layer segments cannot do. This search is the "sublinear plan" in the paper's experiments: every Section 3 optimization plus Algorithm 3's plan.

Recursion: logarithmic memory

Section 4.4 treats a segment as one large operator whose backward pass runs a small graph of its own. Nothing outside the segment depends on how that inner graph is executed, so the same checkpointing can be applied inside it. Let $g(n)$ be the memory to run forward and backward over $n$ layers when $k$ intermediate results are stored and the same strategy is applied to the pieces between them. The $k$ stored results cut the chain into $k+1$ pieces of length $n/(k+1)$, and only one piece is being worked on at a time:

Each level of the recursion adds $k$ and divides the length by $k+1$. Dividing $n$ by $k+1$ until one layer is left takes $\log_{k+1} n$ levels, so

With $k = 1$, each level keeps one checkpoint in the middle of its piece and halves the length, and $g(n) = \log_2 n$. For $n = 4096$ that is 12 stored activations instead of 4,096. The cost is a recomputation per level: about $\log_2 n$ forward passes, or $O(n \log n)$ forward compute in total. The paper says this version "may not be used commonly" and presents it as a demonstration of how far the trade goes.

Figure 6 plots equation (3) as a function of the number of levels $p$, taking $k + 1 = n^{1/p}$ so the recursion has exactly $p$ levels, and reading $p$ as the number of forward passes. At $p = 2$, $k = 63$ and the memory is $2 \cdot 63 = 126$, which is the $\sqrt{n}$ scheme ($2\sqrt{n} = 128$ before rounding the segments). At $p = 12$, $k = 1$ and the memory is 12. Drag $p$ and compare the drop from 1 pass to 2 with the drop from 2 to 12.

Choosing the checkpoints optimally at every level of this recursion is the revolve algorithm of Griewank and Walther (2000), which the paper cites. Their book states the general trade: for a constant $c$, $O(c\,n)$ compute buys $O(n^{1/c})$ memory. $c = 1$ is plain backprop, $c = 2$ is the $\sqrt{n}$ scheme, and taking $c$ up to $\log n$ gives $\log n$ memory.

What the experiments measured

All experiments run on MXNet and report two quantities. The feature-map cost comes from the static allocator's plan: exact, available without running, and limited to activations. The total cost is measured with nvidia-smi (NVIDIA's device-monitoring tool) while training on a Titan X, a 12 GB card, and includes parameters and workspace; plans that ran out of memory have no measurement. The paper compares five plans: no optimization, in-place only, sharing (all of Section 3), drop bn-relu, and the sublinear plan.

The convolutional benchmark is the identity-mapping ResNet of He et al. on 224×224 ImageNet-size inputs at batch size 32, made deeper by adding layers to each stage, from about 100 to about 1,000 layers (a convolution, batch norm, ReLU group counts as one layer). The recurrent benchmark is a four-layer LSTM (a recurrent network that carries a state from one time step to the next) with 1,024 hidden units, batch size 64, a 50-dimensional input per step and a 5,000-way softmax output, a speech-recognition setting, unrolled from about 50 to about 1,600 time steps. An unrolled RNN is a chain with one link per time step, so sequence length plays the role of depth.

On the ResNet feature-map view, every plan except the sublinear one is a straight line of about the same slope as no optimization, shifted down: in-place and sharing are the constant factor of two to three. The sublinear plan grows from about 1.6 GB at 100 layers to about 4.0 GB at 1,000, roughly 2.5 times for 10 times the depth. Switch to total GPU: sharing measures 11.1 GB at about 200 layers and has no point after that, while the sublinear plan reaches 6.6 GB at 1,000 layers, the paper's "less than 7GB".

The abstract's "from 48G to 7G" takes 48 GB from the feature-map view (the sharing plan at 1,000 layers, a plan that could not run) and 7 GB from the total view. Compared within one view, feature maps drop from about 48.6 GB to about 4.0 GB, about 12 times. The roughly 2.6 GB between the sublinear plan's feature-map estimate and its measured total at 1,000 layers is parameters, workspace and other memory outside the feature-map count, which recomputation does not touch.

On the LSTM, in-place matters more than on the ResNet: it lets the weight gradient of each time step be added into one buffer, instead of allocating a gradient buffer at every step. The sublinear plan's advantage grows with length, from about 3 times at 50 steps to about 10 times at 1,600 (2.46 GB against 23.8 GB of feature maps); the paper summarizes it as "more than 4x". Measured, the sharing plan fills the 12 GB card at about 800 steps and has no point beyond, while the sublinear plan runs 1,600 steps in about 3 GB.

For speed, the paper times 20 batches on one Titan X (its Figure 7): the sublinear plan costs about 30% more time per batch than the sharing plan on both networks, and time per batch stays linear in depth for both.

Where it came from, where it went

The paper does not claim the idea as new. Its Section 2 says that dropping intermediate results "is also known as gradient checkpointing technique in automatic differentiation literature", citing Griewank and Walther's revolve, and the acknowledgements thank David Warde-Farley for pointing out that connection. What the paper added was the method for general deep-network graphs (convolutional and recurrent), mirror counts in the graph builder, the budget planner, the combination with buffer sharing, and measurements showing the cost was 30%.

In PyTorch the whole of Algorithm 1 is one call. For a sequential model of 64 blocks, eight segments is the $\sqrt{n}$ choice:

import math, torch
from torch import nn
from torch.utils.checkpoint import checkpoint_sequential

n = 64
net = nn.Sequential(*[nn.Sequential(nn.Linear(1024, 1024), nn.ReLU())
                      for _ in range(n)])
x = torch.randn(32, 1024, requires_grad=True)
k = int(math.sqrt(n))                          # 8 segments of 8 blocks
y = checkpoint_sequential(net, k, x, use_reentrant=False)
y.sum().backward()      # re-runs each segment's forward before its backward

Large Transformer training uses the same switch. Megatron-LM and DeepSpeed expose activation checkpointing per layer, and ZeRO attacks the other half of the memory bill (parameters, gradients and optimizer state), which checkpointing leaves alone. Korthikanti et al.'s selective activation recomputation is Section 4.2's idea applied to Transformers: recompute only the attention operations that take a lot of memory and little compute, instead of whole layers.

Three families of methods now reduce activation memory, and they differ in what they assume. Checkpointing recomputes dropped activations from stored ones and works on any network. Reversible residual networks (Gomez et al., 2017) reconstruct each layer's input from its output, reaching constant memory in depth, but only for layers built to be invertible. Neural ODEs use the adjoint method: they solve a second differential equation backward in time to get the gradients, also in constant memory, with numerical error in the reconstruction as the price. FlashAttention also recomputes in its backward pass, but what it stores is the attention output and the per-row softmax statistics, and what it recomputes is the $N \times N$ attention matrix, block by block in on-chip memory, so that the matrix is never written to GPU main memory.