Activation stashing and rematerialization

Index saved state by stage, microbatch, and relevant parameter version. Even equal-shaped tensors are not interchangeable when they arose from different data or weights. A reusable buffer must remain owned until its backward and any communication readers finish. Pipeline scheduling changes how many bundles overlap; checkpointing changes what each bundle contains. These are separate optimization dimensions. A lower theoretical bundle count does not prove an actual memory reduction if the implementation keeps duplicate references or additional communication storage alive.

Estimate the saved-state peak from a lifetime trace rather than summing every tensor ever allocated. Include metadata and any live checkpoint inputs, then distinguish those bytes from weights, optimizer state, temporary workspace, and allocator reserve. A checkpointed region can still retain its input and outputs needed elsewhere, so its storage is not generally zero. Select regions according to the actual graph and backward requirements. Tensors needed by noncheckpointed branches can remain live even if the region itself does not save their intermediates.

Recomputing a forward region during backward must reproduce the state needed for the intended gradient. Dropout requires appropriate random-state handling; mutable counters, batch-normalization running statistics, external logging, and other side effects can execute differently or twice. A global variable changed between forward and recomputation can silently alter the function. Holding the correct parameter version is also necessary. Do not assume that using the same input shape or calling a checkpoint API proves numeric equivalence to the original execution.

PyTorch 2.8 recommends explicitly selecting the checkpoint implementation. Its nonreentrant variant can stop recomputation once needed values are rebuilt, while the reentrant variant reruns the whole function. Default RNG preservation covers CPU and a selected device type under documented rules; moving tensors to a new device inside the function can defeat the assumed coverage. The default nonreentrant determinism check compares metadata such as shapes, dtypes, and devices rather than proving equality of tensor values. Verify gradients and mutable state separately.

Compare checkpointed and uncheckpointed runs at the same model, data, precision, update boundary, and intended random behavior. Report retained-state bytes and observed peak allocation with their definitions, plus recomputation work and end-to-end time. A memory saving can enable a larger feasible microbatch and thereby improve throughput despite extra arithmetic. Conversely, a region with little saved state and expensive replay can be a poor choice. The decision depends on which constraint limits the complete training configuration.

The cache lab demonstrates only working-set capacity. Reset, set capacity to 4 and stride to 4, predict hits after seventeen requests, then Step once: thirteen hits. Compare a fresh Reset with both capacity at 2 and stride at 4 before Step: zero hits. The exact toy LRU stream starts empty with sixteen requests already counted. It does not rematerialize activations, preserve RNG state, trace autograd, or represent the lifetimes of saved tensors in a real pipeline.