Gradient checkpointing & memory maths

1

Concept

Backpropagation needs values from the forward pass to compute gradients. Autograd therefore retains activations, attention intermediates, and other tensors until their backward functions run. In a deep Transformer, these saved values can consume more memory than the weights. Activation checkpointing, often called gradient checkpointing, reduces that memory by deliberately forgetting and recomputing.

Divide the network into segments. During the original forward, save only each segment’s input or boundary output instead of every internal activation. During backward, rerun the segment’s forward computation from the saved boundary, reconstruct its intermediates, calculate gradients, then release them. The parameters are not rolled back; “checkpoint” here means an activation boundary, not a file on disk.

The tradeoff is direct: lower peak activation memory for extra computation. If every Transformer block is checkpointed, much of each block’s forward work runs again during backward. The exact overhead is not simply double total training time because backward, communication, and non-checkpointed work remain, and recomputation may overlap differently. Measure step time and realized throughput.

A memory budget should separate categories:

Checkpointing primarily attacks $M_{activations}$. It does not shard parameters or Adam moments. It may also not remove every attention temporary, depending on kernel and where the checkpoint boundary sits. A job dominated by optimizer state needs ZeRO/FSDP; one dominated by long-sequence activations may benefit greatly from checkpointing.

Activation size grows with batch, sequence length, model width, and number of saved layers. Ordinary attention can add terms quadratic in sequence length, while memory-efficient attention kernels avoid materializing the full probability matrix. Combine estimates from tensor shapes with measured allocator peaks, because fused kernels and framework bookkeeping alter what is retained.

Stochastic operations require care. If dropout produced one mask during the original forward and another during recomputation, backward would differentiate a different function. Framework checkpoint utilities preserve or restore random-number state, with a performance cost. Stateful modules, side effects, mutable caches, and data-dependent external calls can also make recomputation incorrect.

Granularity changes the tradeoff. Large segments save more boundaries but recompute more work and may retain large segment inputs. Small segments keep more checkpoints and invoke more framework overhead. Selective policies can checkpoint attention and FFN branches differently. Profile memory and compute instead of choosing “every N layers” by folklore.

Pipeline parallelism complicates scheduling because recomputation competes with other microbatches. FSDP may all-gather weights again during recompute unless parameters remain materialized, trading communication against memory. Compilation systems may fuse or reorder regions. The effective unit is the whole distributed schedule, not an isolated function.

A calculator should let learners enter batch, $T$, $d_{model}$, layers, dtype bytes, FFN expansion, and checkpoint interval, then show an explicitly simplified estimate. It must label omitted categories and avoid presenting theoretical tensor counts as an allocator guarantee.

The durable mental model is time–space exchange. Save enough boundary state to reconstruct the forward graph; discard bulky intermediates; pay compute to recreate them just before their gradients are needed. The right decision follows from a measured memory breakdown, not the generic fact that the model is large.