What correctness issues should you consider when using activation checkpointing?

How to think about it

Activation checkpointing is correct only when replaying a checkpointed region during backward reconstructs the same values and autograd dependencies as the original forward pass. The main hazards are side effects, random operations, changing control flow or devices, and choosing an incompatible reentrant mode.

Why replay creates a correctness problem

Autograd, PyTorch’s automatic differentiation system, normally saves intermediate tensors from the forward pass so it can use them when calculating gradients during backward. Activation checkpointing saves fewer of those tensors. It keeps enough information to call the region again later, then recomputes the missing intermediates.

That saves memory at the cost of extra computation. For example, one float16 activation with shape [8, 2048, 4096] occupies:

8 × 2048 × 4096 × 2 = 134,217,728 bytes, or 128 MiB.

A large network can hold many such tensors. Checkpointing trades some of that memory for replayed matrix multiplications and other forward operations.

The correctness requirement is simple to state: if the original region computes f(x, state), replay must compute the same relevant function from the same x and equivalent state. “Relevant” matters because backward may need only some intermediate tensors, not necessarily the final output again.

If replay takes a different branch, uses a different dropout mask, sees a changed buffer, or produces a different shape, the gradient may be different or PyTorch may raise an error. The dangerous case is a numerically different result with the same shape and dtype. That can pass through the program without an obvious exception.

Randomness: dropout is handled, but not magically

Dropout is the usual example. During the original forward pass, dropout samples a mask. During replay, backward needs the same mask. A different mask means the gradient belongs to a different stochastic computation.

PyTorch checkpointing preserves the PyTorch random-number-generator state by default. It saves the CPU state and the state for one inferred non-CPU device type, prioritising CUDA when appropriate. This normally makes operations such as torch.nn.functional.dropout replay with the same random values.

That guarantee has boundaries:

• Setting preserve_rng_state=False can reduce overhead, but a stochastic region is no longer equivalent to the uncheckpointed version.
• Python’s random module and NumPy’s random generator are not controlled by PyTorch’s RNG preservation.
• If a region moves tensors to a new device internally, checkpointing may not know that device’s RNG state must be saved.
• A nondeterministic kernel or external state can still produce different results even when the RNG state matches.

A common mistake is to seed the program once and assume that every later execution is therefore identical. A seed determines a sequence; it does not make two executions identical if one execution consumes random numbers in a different order.

Side effects can happen twice, or not at all

A checkpointed function should behave as close to a pure function as practical: its result should come from its inputs and stable state, rather than from mutations performed while it runs.

Consider a convolutional model with BatchNorm inside a checkpointed region. In training mode, BatchNorm updates its running mean. With the default momentum of 0.1, suppose the old running mean is 0 and the current batch mean is 10.

The original forward update gives:

0.9 × 0 + 0.1 × 10 = 1

If replay performs the update again, the running mean becomes:

0.9 × 1 + 0.1 × 10 = 1.9

The current batch’s output may still look reasonable, so the bug can hide. Later batches and evaluation use the wrong running statistics.

The same problem appears with counters, cache writes, metrics, list appends, custom logging, and in-place mutation of buffers or inputs. A side effect before a needed intermediate may happen twice. With non-reentrant checkpointing, replay can also stop early once it has reconstructed everything backward needs, so a side effect later in the region may not happen during replay at all. Either behaviour is a poor contract to build on.

Keep stateful operations outside the checkpointed region when possible. If they must remain, test the state itself, not only the loss and gradients.

Control flow, devices, and context must remain stable

Data-dependent control flow is not automatically wrong. A branch based only on unchanged input values will normally take the same path during replay.

The problem is control flow that depends on mutable or untracked state. Examples include a global flag, a counter, a changing module.training value, a Python random draw, or a tensor captured from a closure and modified between forward and backward.

Device and dtype changes create similar problems. Moving a tensor between CPU and GPU inside the region, changing autocast settings, or switching a module from training to evaluation mode can make replay differ from the original. A shape, dtype, or device mismatch may be caught by PyTorch’s checkpoint checks. A same-shaped but numerically different tensor may not be.

The default determinism check is therefore useful but limited: it checks metadata such as shape, dtype, and device, not that every recomputed value is numerically equal. Treat it as a smoke alarm, not as a proof of gradient correctness.

use_reentrant is a correctness choice

Pass use_reentrant explicitly. The two modes are not interchangeable implementation details.

The recommended non-reentrant mode, use_reentrant=False, records the autograd graph during the original forward. It supports the usual backward APIs, including torch.autograd.grad, handles nested input structures more fully, and can stop replay early when the required intermediates have been reconstructed.

Reentrant mode, use_reentrant=True, runs the original forward under torch.no_grad() and rebuilds the graph during backward through a custom autograd boundary. It has stricter requirements: at least one input and output must participate in gradient computation, and detached tensors or some nested structures can fail even when the uncheckpointed code works.

That means a region using higher-order gradients, detached auxiliary outputs, or unusual input containers should generally use non-reentrant mode and be tested explicitly. Older code that relied on the historical default deserves particular suspicion.

A concrete dropout check

This small test compares a normal forward with a checkpointed forward using the same initial weights, inputs, and RNG state. With RNG preservation enabled, the losses and gradients should agree within the chosen tolerance.

import torch
import torch.nn.functional as F

from torch.utils.checkpoint import checkpoint

torch.manual_seed(7)
x0 = torch.randn(32, 1024)
w0 = torch.randn(1024, 1024)
initial_rng = torch.get_rng_state()

def evaluate(use_checkpoint):
    x = x0.detach().clone().requires_grad_()
    w = w0.detach().clone().requires_grad_()

    def region(x, w):
        z = F.dropout(x, p=0.1, training=True)
        return (z @ w).square().mean()

    if use_checkpoint:
        loss = checkpoint(
            region,
            x,
            w,
            use_reentrant=False,
            preserve_rng_state=True,
        )
    else:
        loss = region(x, w)

    loss.backward()
    return loss.detach(), x.grad.detach(), w.grad.detach()

torch.set_rng_state(initial_rng)
reference = evaluate(False)

torch.set_rng_state(initial_rng)
replayed = evaluate(True)

for expected, actual in zip(reference, replayed):
    torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-6)

If this comparison fails after adding checkpointing, first inspect random-state handling, mutable state, autocast, and the checkpoint boundary. Do not immediately loosen the tolerance. A tolerance can hide a replay bug just as effectively as it can accommodate harmless floating-point noise.

The production pattern

I would introduce checkpointing in this order:

1. Isolate a mostly stateless region and pass its meaningful tensor inputs explicitly.
2. Choose use_reentrant=False or True deliberately rather than relying on a default.
3. Keep RNG preservation enabled for stochastic operations unless a measured design decision says otherwise.
4. Compare the uncheckpointed and checkpointed versions using the same inputs, weights, mode, and RNG state.
5. Compare the loss, input gradients, parameter gradients, and, where relevant, the resulting optimizer update.
6. Test training and evaluation, mixed precision, dynamic shapes, and every device type used in production.
7. Benchmark the real model shape. Checkpointing saves memory, but replay increases compute and can add latency.

Do not use it as an ordinary inference optimisation. Without a backward pass, there are usually no training activations to reclaim. Also avoid wrapping code with unavoidable side effects. Split the pure tensor computation from the state update instead.

What they’ll ask next

Does checkpointing make dropout deterministic?
Usually, for PyTorch RNG operations, yes: PyTorch saves and restores the relevant RNG state by default. It does not control Python or NumPy randomness, and disabling preservation removes equivalence with the uncheckpointed computation.

Why prefer non-reentrant checkpointing?
It records the forward autograd graph, supports more backward patterns and input structures, and can stop replay early. Reentrant mode is more restrictive and is mainly appropriate when its constraints are understood and tested.

How would you verify correctness?
Run both implementations from identical inputs, parameters, modes, and RNG states. Compare outputs and gradients with appropriate tolerances, then compare an optimizer step. Include tests for stateful modules and stochastic paths; a successful loss calculation alone is not enough.

Say this in the interview

“Checkpointing is correct only when replay sees the same inputs, state, RNG behaviour, and execution path, so I use a pure region, choose the reentrant mode explicitly, preserve RNG for stochastic layers, and compare gradients with an uncheckpointed baseline.”