Memory Budget

## What actually fills memory during a forward pass When a network runs forward, each layer produces an output tensor, and reverse-mode automatic differentiation needs many of those intermediate values a second time to form local gradients on the way back. The backward pass walks the graph in reverse, so the values produced earliest are needed last. By default they all sit in device memory from the moment they are produced until the backward pass consumes them. Peak activation memory therefore grows with depth and with everything that multiplies the size of a single intermediate: batch size, sequence length, image resolution, channel width. This is the term that usually pushes a training job over the device limit, because unlike the weights it scales with the batch. ## The idea Activation checkpointing — also called gradient checkpointing or rematerialization — refuses to hold that whole stack. The network is cut into segments. Only the tensor entering each segment is retained; everything computed inside a segment is released as soon as the segment's output has been produced. When the backward pass reaches a segment, the segment is executed forward a second time from its retained input, the interior values are reconstructed, the gradients for that segment are computed from them, and the reconstructions are dropped again before moving to the next segment. At any instant only the small set of boundary tensors plus one segment's interior is alive. The technique is trading a resource that is elastic (compute time) for one that is a hard ceiling (device memory). Running slower is a decision; running out of memory is a crash. ## Counting the cost A good accounting rule: if a forward pass costs F, the backward pass costs roughly 2F, because at most layers it forms both a gradient with respect to the layer's input and a gradient with respect to its parameters. An ordinary step is therefore about 3F. Checkpointing the whole network adds one more forward, giving about 4F — roughly a third more compute per step. This is why the technique is usually quoted as costing about thirty percent of step time when everything is checkpointed; checkpointing only part of the network costs proportionally less. Wall-clock time often moves by less than the compute ratio, for a practical reason: teams rarely enable checkpointing and keep the same batch size. The freed memory is spent on a larger batch, which usually improves device efficiency, so some of the recompute cost is earned back as better utilization. ## What it does not help Checkpointing touches the activation term only. Parameters, their gradient buffers and optimizer state are untouched, because none of them is a recomputable intermediate — they persist across the whole step by definition. If a model cannot fit with a batch of one and no activations at all, checkpointing cannot rescue it; that situation needs a different lever such as sharding or reducing the model itself. It also does not reduce the number of parameters, change the learning rate that is appropriate, or alter what the optimizer does. ## Correctness Recomputation reproduces the same values from the same inputs, so the gradients handed to the optimizer are the ones an un-checkpointed run would produce. The exception is any layer whose forward pass is not a pure function of its input — a layer that draws random values, or one that mutates internal state as a side effect of running forward. Those need explicit care so the second forward reproduces the first; otherwise the recomputed values do not match the ones that produced the loss. ## A naming trap worth pre-empting Activation checkpointing has nothing to do with saving model checkpoints to disk for fault tolerance or for resuming a run. The word is shared, the mechanism is not: one is about which intermediate tensors stay in device memory inside a single step, the other is about periodically persisting parameters. Interviewers ask this deliberately, because a candidate who has only read the flag name often conflates the two. ## When to reach for it Use it when activations dominate the memory profile and the model is deep enough that a segment interior is much smaller than the whole stack — long-sequence encoders, high-resolution vision models, anything where the useful batch size is being squeezed to one or two. Skip it when the step is already compute-bound and memory is comfortable, because you would be paying a third of your throughput for room you do not need.