Gradient Checkpointing Math

Why activations, not weights, dominate training memory

Inference only needs the current layer’s output, so its memory footprint is dominated by the weights. Training is different: the backward pass computes ∂L/∂W for each layer, and the chain rule needs that layer’s input activation to form the gradient. So every intermediate result produced in the forward pass must be kept alive until the backward pass consumes it.

For a network of L layers, that means holding roughly L activation tensors at once. For a transformer, one layer’s stored activations scale as b · s · h (batch, sequence length, hidden width), plus attention and MLP intermediates. With modern depths and sequence lengths this term grows large and linearly in L, while the weights are fixed regardless of batch or sequence. That is why a model that fits for inference can blow past memory in training: it is the activations, scaling as O(L · b · s · h), that overflow.

The core idea: drop activations, recompute them later

Checkpointing exploits a simple asymmetry. An activation is cheap to recompute (you re-run part of the forward pass) but expensive to store (it occupies memory for the whole forward-then-backward window). If memory is the binding constraint and you have spare compute, trade one for the other.

Concretely: designate a subset of layers as checkpoints and store only their activations during the forward pass; discard everything in between. When the backward pass reaches a discarded region, re-run the forward computation for that region — starting from the nearest stored checkpoint — to regenerate the activations it needs, use them for the gradient, then free them again. You pay for the segment’s activations only transiently, while backpropagating through it, instead of for the entire step. The whole design question is then: how many checkpoints, and where?

The two naive extremes bound the trade

Before optimizing, pin the endpoints. Store everything (the default): memory is O(L), compute is one forward plus one backward — no recomputation, no overhead. This is fast but memory-hungry.

Store nothing but the inputs: memory for stored activations drops toward O(1), but to get the gradient at the last layer you must recompute the forward pass up to that point, and again for the layer before it, and so on. Done naively that is O(L^2) recomputation — quadratic, and disastrous for deep nets. So the extremes are O(L) memory / 1× compute at one end and O(1) memory / O(L^2) compute at the other. Neither is what we want. The interesting regime is in between, and the right structure — segmenting the network and storing one checkpoint per segment — avoids the quadratic blow-up entirely.

Segmenting the network: the memory model

Divide the L layers into segments of length k, giving L/k segments. During the forward pass, store only the activation at each segment boundary — that is L/k stored checkpoints. During the backward pass, process one segment at a time: starting from that segment’s stored input checkpoint, recompute its k internal activations, backpropagate through them, then release them before moving to the previous segment.

So peak activation memory has two parts:

M(k) = (L / k)   +   k
        checkpoints    recompute buffer
               (one live segment)

The first term is the sparse checkpoints held for the whole step; the second is the transient buffer of one segment’s worth of activations, live only while that segment is being reconstructed. Everything is in units of one layer’s activation tensor. Minimizing this sum over k is the whole game.

Deriving the sqrt(L) optimum

Minimize M(k) = L/k + k by differentiating with respect to k and setting it to zero:

dM/dk = -L / k^2 + 1 = 0
   =>   k^2 = L
   =>   k* = √L        (optimal segment length)

M(k*) = L/√L + √L = √L + √L = 2√L

The optimum is a segment length of √L, which also means about √L checkpoints. The two terms balance exactly — a classic AM–GM result, where a sum a/x + x is smallest when both pieces are equal. Peak activation memory falls to 2√L = O(√L). For any L larger than a handful of layers this is a large win: the memory needed to backprop through the network now grows with the square root of depth, not linearly. Doubling depth costs only a factor of √2 ≈ 1.41 in activation memory instead of 2.

Where the ~33% recompute overhead comes from

The compute cost is easiest to see in units of a single forward pass. Backpropagation is roughly twice the work of a forward pass — each layer computes a gradient with respect to its input and with respect to its weights — so a normal training step costs about:

normal step  = forward (1) + backward (2)          = 3 units
checkpointed = forward (1) + recompute (1) + backward (2) = 4 units

With checkpointing, every layer’s forward is evaluated at most twice: once in the initial forward pass, and once when its segment is recomputed during backprop. That adds exactly one extra forward pass over the network — one unit on top of three — for an overhead of 4/3 - 1 ≈ 33%. This is the number quoted in practice: gradient checkpointing costs roughly a third more compute per step. Crucially the overhead is constant, independent of k or L, because each activation is recomputed at most once — that is precisely what segmenting buys over the naive O(L^2) scheme.

A worked example

Take a transformer with L = 64 layers. Storing every layer’s activations costs 64 units. Choose segment length k = √64 = 8: you store 64/8 = 8 boundary checkpoints and hold at most 8 layers’ activations live while recomputing a segment, for a peak of 8 + 8 = 16 units.

That is a 64 → 16, or 4×, reduction in activation memory, paid for with about 33% more compute. Put in tensor terms: if one layer’s activations are, say, b·s·h = 1 × 4096 × 4096 ≈ 16.8M elements at 2 bytes each ≈ 34 MB, then full storage is 64 × 34 ≈ 2.1 GB while checkpointed storage is 16 × 34 ≈ 0.55 GB. Freeing 1.5 GB is often the difference between a batch size that fits and an out-of-memory crash — and the 33% slowdown is a bargain when the alternative is not training at all.

Segment selection when layers are not equal

The clean √L result assumes every layer costs the same to store and to recompute — true for a stack of identical transformer blocks, which is why uniform spacing is the sensible default. When layers are heterogeneous, the objective generalizes: minimize total recompute FLOPs subject to a memory budget, or minimize memory subject to a FLOP budget. This is a knapsack-flavored problem where each candidate checkpoint has a memory cost (the activation it pins) and a recompute benefit (the work it saves re-deriving downstream).

The practical heuristic: place checkpoints so that segments are balanced in recompute cost, not merely in layer count. Prefer to store activations that are expensive to recreate (a big attention score matrix) and to discard those that are cheap to recreate (an elementwise activation or a normalization). That insight is the seed of selective recomputation, where you checkpoint at a finer, op-level granularity rather than only at whole-layer boundaries.