Checkpointing trades saved activations for compatible recomputation

Consider a chain of L similar layers divided into segments of s layers. A simple schedule keeps segment boundaries and, during reverse work on one segment, recomputes that segment's internal activations. An idealized payload budget has about L/s boundary slots plus s local slots rather than L saved layer outputs. This gives a useful square-root-style tradeoff for a uniform chain, not a universal optimal policy. Real graphs have branches, nonuniform activation sizes, different recomputation costs, and kernels whose workspaces can dominate the activation payload.

The original sublinear-memory research analyzes and demonstrates particular checkpoint schedules; its asymptotic result depends on the model and implementation assumptions. Saving more boundaries can reduce replay while retaining more memory. Nested or optimized schedules make different tradeoffs, and offloading introduces bandwidth and synchronization costs. Include parameters, gradients, allocator reserves, input/output storage, communication buffers, and temporary workspaces when measuring total peak. A theoretical count of activation slots should be labeled as such rather than advertised as an exact accelerator-memory guarantee.

Random operations require a compatible RNG-state policy if replayed activations must match. PyTorch 2.9 documents checkpoint RNG preservation for CPU and one additional device type inferred from the tensor arguments, with limitations when several device types or movement to a new device are involved. It cannot anticipate arbitrary device changes inside the function. PyTorch 2.9 documents that torch.compile preserves RNG regardless of the preserve_rng_state flag; on execution paths that honor disabling it, the reproducibility assumption changes. Other random generators or external state also need a contract; a fixed global seed alone does not establish the exact same draw sequence after a different execution schedule.

A checkpointed function that updates a global counter, appends output, performs an external side effect, or changes mutable model state can behave differently when recomputed. Some implementations stop replay once needed residuals are reconstructed, while others rerun a whole region; PyTorch 2.9 distinguishes reentrant and non-reentrant variants and their transform restrictions. Select and test the documented variant explicitly. Do not rely on replay executing an exact number of side effects. The safe teaching model is a compatible computation whose necessary values can be regenerated without changing its intended semantics.

Compare gradients with a trusted non-checkpointed reference at fixed inputs and state, then measure peak allocated and reserved memory, step time, extra kernels, transfers, and variation. Include higher-order or custom-rule paths if the application needs them. A partial replay can make simple layer-count estimates inaccurate in either direction. Verify that a saved boundary stays live only as long as intended and that the chosen checkpoint policy does not retain another reference to the supposedly discarded activation. The goal is a correct computation under a useful resource budget, not merely a lower number in one memory category.

The cache lab illustrates finite remembrance for a hypothetical pure-value lookup, not an AD schedule. Reset, set BOTH C=4 and s=2, predict zero hits among the N=16 requests at t=0, then Play or Step. Reset and re-set BOTH C=8 and s=2; predict eight hits. N=min(4096,16+floor(2t)); edits retain t/recompute from empty. A miss would require a separately assumed fetch or recomputation in an application, but the lab does neither. It has no activation graph, checkpoint boundaries, RNG, device memory, or measured replay time.