Activation Checkpointing › Recomputation Correctness and Determinism
Language AI HandbookAll chapters Activation Checkpointing › Recomputation Correctness and DeterminismOne section · Sections in this chapter10 / 13 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
A subtle but important requirement for activation checkpointing is that the recomputed activations must be numerically identical to the activations produced during the original forward pass. If the recomputation produces different values, the gradients computed during backpropagation will be incorrect, leading to silent training failures that can be very difficult to diagnose.
For most operations in standard transformer architectures, recomputation is deterministic: given the same input tensor, the same layer produces the same output. However, several operations can introduce non-determinism:
Dropout. Dropout randomly zeroes a fraction of tensor elements during training. A naive recomputation would re-apply dropout with a different random mask, producing different activations and incorrect gradients. The correct behavior is to apply the same mask during recomputation as was applied during the original forward pass. PyTorch handles this automatically by saving the random number generator state at each checkpoint boundary and restoring it during recomputation.
Random number generation in custom layers. Any custom operation that draws from a random distribution during the forward pass must either be made deterministic or must save its random state alongside the checkpoint tensor. PyTorch's non-reentrant checkpointing saves the CPU and CUDA random states automatically, which covers most use cases.
Non-deterministic CUDA operations. Some CUDA operations (particularly certain reduction operations) are non-deterministic by default on NVIDIA GPUs, meaning they can produce slightly different results on different runs due to floating-point non-associativity. This rarely causes correctness issues for checkpointing, but can complicate debugging if you are comparing activations between the original forward pass and the recomputed version.
The PyTorch implementation handles the random state correctly by design. When you call checkpoint.checkpoint with use_reentrant=False, the implementation records the RNG state before the forward pass and restores it before the recomputation, making sure dropout and similar operations produce identical results.