Activation Checkpointing › How PyTorch Implements Activation Checkpointing

Understanding the implementation mechanics helps you debug issues and use the API correctly. The PyTorch implementation of checkpoint.checkpoint with use_reentrant=False works as follows.

The non-reentrant implementation uses a custom torch.autograd.Function subclass. This class has two class methods, forward and backward, that define the computation during both passes. The forward method runs the user-provided function (the module to be checkpointed) but uses torch.no_grad() to prevent PyTorch from building an autograd graph during this run. This is why activations are not retained: with no_grad() active, PyTorch discards the computation graph as it is built.

The forward method saves the inputs to the checkpointed region (the boundary tensors), the function itself, and the RNG state. It does not save any intermediate activations.

During the backward method, PyTorch calls the saved function again on the saved inputs, this time with autograd enabled. This recomputation run builds a fresh computation graph for the checkpointed segment. The backward pass then traverses this fresh graph to compute gradients, after which the graph is immediately discarded.

The most subtle aspect of the implementation is RNG state management. Suppose your checkpointed segment contains a dropout layer. Dropout samples a random mask during the forward pass. If the backward pass recomputes dropout with a different random mask, the gradients will be incorrect (they will reflect a different network configuration than the one that produced the forward pass output).

PyTorch addresses this by recording the RNG state on both CPU and CUDA immediately before the checkpointed region executes during the forward pass. The state includes the full state of the PyTorch random number generator (a Mersenne Twister on CPU and a device-specific generator on CUDA). During the recomputation in the backward pass, PyTorch restores this state before re-running the function, making sure that any random operations produce exactly the same results as the original forward pass.

This RNG state save-and-restore happens automatically and transparently. You do not need to do anything special when using dropout or other random operations inside a checkpointed region.

The preserve_rng_state argument to checkpoint.checkpoint (defaulting to True) controls whether the RNG state is saved and restored. You might set this to False if your checkpointed region contains no random operations and you want to avoid the overhead of saving the RNG state. In practice, the overhead is negligible (saving a few hundred bytes), so leaving it at True is safe.

When checkpoint boundary tensors are saved, they are saved as regular CUDA tensors in HBM. The PyTorch garbage collector will not free them until they are no longer referenced by any Python object or autograd graph. This is what ensures they survive until the backward pass needs them for recomputation.

One practical implication: the checkpoint boundary tensors are included in PyTorch's memory tracking. When you query torch.cuda.memory_allocated(), you will see the boundary tensors contributing to the reported allocation. This is correct behavior: they are truly allocated. The memory reduction from checkpointing shows up as a reduction in the total allocated memory, not a reduction in what is tracked.

A subtlety in the non-reentrant implementation: during the recomputation run in the backward pass, PyTorch builds a complete autograd graph for the checkpointed segment. This graph is then used to compute gradients within the segment. Once the segment's gradients have been computed, the graph is freed. This means the memory spike during recomputation includes the intermediate activations plus the autograd graph nodes, which add a small constant overhead per operation. For large segments with many operations, this overhead can be noticeable.

The reentrant implementation handled this differently: it re-entered the Python autograd machinery directly, which avoided building an explicit graph but had other correctness issues. The non-reentrant approach trades slightly higher peak memory during recomputation for correctness and compatibility with modern PyTorch features.