Checkpointing

Checkpointing is the practice of saving intermediate state so that work can be resumed or partially recomputed instead of repeated from scratch, and in machine learning the word names two quite different techniques that are often confused. Training checkpointing writes the model's weights, optimizer state and progress counters to disk so a long run can survive crashes. Activation checkpointing, also called gradient checkpointing or rematerialization, deliberately discards some intermediate values during the forward pass and recomputes them during the backward pass, trading extra compute for lower memory. They share a name and the same underlying trade of storage against recomputation, but they solve different problems and live in different parts of the system.

Training checkpoints exist for fault tolerance and for flexibility. A checkpoint must contain more than the weights: the optimizer state (for Adam, two extra values per parameter), the learning-rate scheduler position, the step count, the random number generator states and the data loader position are all needed to resume a run that behaves as though it never stopped. For a large model this can mean hundreds of gigabytes to terabytes per save. Clusters with thousands of accelerators fail often; Meta's 2024 report on training Llama 3 documented hundreds of interruptions during a roughly 54-day training period, most from hardware faults, so frequent saves bound how much work a failure destroys. Checkpoints also support picking the best epoch by validation score, branching experiments, and later fine-tuning from an intermediate state.

Activation checkpointing attacks the memory wall of backpropagation. The backward pass needs the forward activations of every layer, so memory grows linearly with depth and often dominates the budget of a large transformer. Chen, Xu, Zhang and Guestrin, in "Training Deep Nets with Sublinear Memory Cost" (2016), showed that storing activations only at about √n well-chosen points of an n-layer network and recomputing the rest in segments reduces activation memory to roughly O(√n) at the price of about one additional forward pass, on the order of a third more compute for a typical training step. The idea descends from older work on reverse-mode differentiation by Griewank and Walther, whose revolve algorithm schedules checkpoints optimally. Frameworks expose it directly: torch.utils.checkpoint in PyTorch and jax.checkpoint, also called remat, in JAX.

The memory savings are real, but the technique is a tool, not a free lunch. The recomputation cost is why practitioners apply it selectively; Korthikanti and coauthors (2022) described selective activation recomputation for large transformers, which re-derives only the cheap but memory-hungry pieces such as attention scores while keeping the expensive matrix multiplications stored. It is also incompatible with naive handling of randomness: dropout inside a checkpointed segment must reuse the same random mask on recomputation, or gradients are computed against a different forward pass than the one that produced the loss. Pairing it with FlashAttention, which also recomputes rather than stores the attention matrix, shows the same philosophy at a finer grain.

Common pitfalls follow the two meanings. A training checkpoint that saves only weights cannot resume a run faithfully, because losing the optimizer state silently resets the adaptive statistics. A save interrupted midway can leave a corrupted file, so robust systems write to a temporary name and then rename atomically, or keep several recent versions. For activation checkpointing the usual error is expecting a speedup: it slows each step to buy room for a larger batch or longer context. Finally, the word checkpoint is also used for released model weights, as in a downloadable pretrained checkpoint, which is a third, purely distributional sense.

Read checkpointing as a deliberate trade between remembering and redoing: save state to disk so a crash costs minutes rather than weeks, or discard and rebuild activations so a model that would not fit in memory can be trained at all, and keep clear which of the two is meant.