Why architecture matters here
Architecture matters here because activation memory is the wall that most training runs hit first, and checkpointing is the cheapest lever for moving it. Unlike buying more or bigger accelerators, or sharding a model across devices with all the communication cost that entails, checkpointing is a local, per-device technique that requires no extra hardware and no change to the model's math. It simply reorganizes when activations exist in memory, and that reorganization can free the majority of the activation footprint.
The problem is genuinely structural, not incidental. Backpropagation is defined by the chain rule: to compute the gradient of the loss with respect to a layer's inputs, you need that layer's inputs and often its outputs. Those values were produced on the forward pass, potentially thousands of operations earlier, and by default the framework's autograd engine holds a reference to every one of them so it can walk the graph backward. The retained graph is the memory cost, and it grows with depth because every layer adds another set of tensors that must survive until its backward step.
Consider the arithmetic to feel the pressure. Suppose each transformer layer retains activations proportional to sequence length times hidden size times batch, and you have N layers. Naive training holds all N layers' worth simultaneously at the peak, so memory is O(N) in depth. If instead you checkpoint every layer — keeping only each layer's input and recomputing its internals on the way back — the peak drops to roughly the memory of a single layer plus the thin set of boundary tensors, at the cost of recomputing every layer's forward once. Segmenting into sqrt(N) groups gives the classic O(sqrt(N)) memory for O(1) extra forward passes, a sweet spot that keeps both memory and compute overhead modest.
The payoff compounds with the other things you want to scale. Longer context windows multiply activation memory linearly, so checkpointing is often what makes 32k or 128k training fit. Larger batches improve gradient quality and hardware utilization, and the memory they free up by checkpointing can be spent on more batch. And because the technique is orthogonal to mixed precision, optimizer sharding, and tensor parallelism, it stacks with all of them — checkpointing handles activations while those techniques handle weights, gradients, and optimizer states.
It is worth being clear that this is a deliberate, tunable trade rather than a free win. Every segment you checkpoint costs a recompute during backprop, so the more aggressively you checkpoint, the slower each step becomes. The engineering judgment is to checkpoint exactly enough to fit the run comfortably and no more, because compute you spend on rematerialization is compute you are not spending on making progress. The architecture gives you a dial; using it well means reading the memory and throughput numbers and turning the dial only as far as the memory constraint actually requires.
The architecture: every piece explained
Top row: the forward pass with checkpointing enabled. Layers run in order as usual, but the module is wrapped so that at chosen checkpoint marks — segment boundaries — the framework saves the inputs to that segment and nothing else. Everything computed inside the segment (the attention scores, the intermediate feed-forward activations, the per-layer outputs that are not boundaries) is allowed to be freed the moment the forward computation moves past it. The result is a small saved-tensor set living on the device: just the handful of boundary activations rather than the full dense graph.
The key mechanism is that the checkpoint wrapper detaches the segment from the autograd graph on the forward pass. Instead of recording every operation inside the segment for later differentiation, it records a single node that says, in effect, 'to get my gradient, re-run this function.' That is why the internal activations can be dropped — autograd is not holding references to them, because it plans to regenerate them rather than remember them.
Middle row: the backward pass. When the gradient flow reaches a checkpointed segment, the framework recomputes the segment: it takes the saved input, re-runs the segment's forward pass to regenerate the internal activations, and only then runs the normal backward step over that freshly materialized graph to produce the local gradients. Those gradients flow onward to the previous segment, whose activations are in turn recomputed, and so on backward through the network. Each segment's internal graph exists only briefly, during its own backward step, and is freed immediately after — so peak memory is bounded by the largest single segment, not the whole network.
Bottom row: the details that make recomputation correct. The segment policy decides how many checkpoints you place and where; fewer, larger segments save more memory but recompute more per backward step. Critically, the recomputed forward pass must be numerically identical to the original, and that is not automatic when the model contains stochastic operations. Dropout and any RNG-driven op must replay the exact same random state on recompute, or the recomputed activations will differ from the ones used to produce the loss, silently corrupting the gradient. The wrapper therefore captures and restores the random-number generator state around each segment so the second forward is a faithful replay of the first.
The ops strip is the feedback loop. Peak memory tells you how much headroom checkpointing bought and whether you can afford a larger batch or longer sequence. Recompute overhead — the fraction of step time spent regenerating activations — tells you the price you are paying. Segment count is the dial connecting the two. Watching all three together lets you land on a policy that fits memory with the least possible throughput cost, rather than blindly checkpointing everything and eating a needless slowdown.
End-to-end flow
Trace a single training step through a checkpointed model. The batch enters the first segment. The framework saves the segment's input activation, marks the segment as a recompute node, and runs its layers forward. As the forward pass moves into the second segment, the internal activations of the first — its attention matrices, its feed-forward hidden states — are no longer referenced by autograd, so the memory allocator is free to reuse them. This repeats segment by segment until the forward pass reaches the loss. At the peak, device memory holds only the boundary activations plus whatever the current segment is actively computing.
Now the loss is computed and the backward pass begins. Gradient flows into the last segment. Before autograd can differentiate that segment, it needs the internal activations that were dropped — so the checkpoint node fires: it restores the saved RNG state for that segment, re-runs the segment's forward pass from its saved input to regenerate the internal activations, and hands the resulting small graph to autograd. Autograd computes the local gradients, the recomputed graph is freed, and the gradient with respect to the segment's input flows to the previous segment.
That previous segment now repeats the ritual: restore its RNG state, recompute its forward from its saved input, differentiate, free. Because only one segment's internal graph is alive at any instant during the backward pass, peak memory during backprop is also bounded by a single segment rather than the whole model. The gradients accumulate into the parameter gradient buffers exactly as they would without checkpointing — the math is identical; only the memory schedule changed.
The RNG replay is what makes this identity hold. Suppose the segment contains dropout. On the original forward pass, dropout zeroed a particular random subset of activations, and the loss was computed with that mask. If the recompute drew a different mask, the recomputed activations would not match the ones the loss actually depended on, and the gradient would be computed against a graph that never produced the loss — a subtle, silent bug that degrades training. By snapshotting and restoring the generator state around each segment, the recompute reproduces the exact same mask, and the gradient is correct.
Step back and count what the schedule bought across that step. During the forward pass, only boundary activations survived, so peak forward memory fell from all-N-layers to boundaries-plus-one-segment. During the backward pass, segments were rematerialized one at a time and freed immediately, so peak backward memory was similarly bounded. The cost was exactly one extra forward pass over the checkpointed regions, paid during backprop. And because RNG state was replayed faithfully, the gradients are bit-for-bit what non-checkpointed training would have produced. Every one of those properties falls out of a single decision: save segment inputs, drop the internals, and recompute on demand.