Activation checkpointing saves some activation storage by recomputing it later, but the replay must match the original forward computation. If a dropout-like segment draws another mask, backward can combine the original loss’s upstream derivative with a different path’s local derivative. An unchanged forward loss does not prove unchanged gradients.

A scalar example below gives the correct gradient 66 and an incorrect replay gradient 0, then separates the two jobs of RNG-state preservation. The code is a Python standard-library teaching model, not a PyTorch or GPU training experiment. 中文版.

This checkpoint is not a saved model file

Activation checkpointing concerns intermediate tensors inside a training step. Backward normally uses selected values from forward. A checkpointed region retains less of that material and recomputes what is needed later. That differs from writing weights and optimizer state to disk for training recovery.

The PyTorch 2.14 reference describes a compute-memory tradeoff and warns that differing forward and backward function behavior, such as a changed global variable, can cause errors or silently incorrect gradients. This article examines that equivalence requirement rather than claiming a particular memory-saving percentage.

Calling a function again is not sufficient on its own. Inputs, parameters, random state and branch-controlling context can all determine the function actually evaluated. Reconstruct the computation matching the original forward, not merely a tensor with the expected shape. A shape check can be useful without establishing that the values or derivatives belong to the same path.

One mask changes the gradient from 66 to zero

Let trainable scalar w=2, input x=3 and target t=1. A simplified dropout multiplier m is either 2 or 0 with equal probability. Forward computes y=w*x*m and loss L=0.5*(y−t)^2. Fix an original draw of m=2: y=12 and L=60.5.

The upstream derivative entering the segment is dL/dy=y−t=11. The matching local derivative is dy/dw=x*m=6, so dL/dw=66. This differentiates one fixed random realization; it is not a derivation of the expected loss.

If recomputation draws m′=0, its local derivative becomes zero. Combining it with the original upstream derivative 11 gives zero. That is not the gradient of the original m=2 path. A newly sampled path can be valid as a separate forward evaluation without being valid material for this backward calculation.

RNG preservation has two separate jobs

First comes local replay. Save state s0 before the original segment, then temporarily restore it during recomputation to reproduce that draw. Setting a seed only once at training startup does not reposition the generator correctly for every later replay. The position in the stream matters as well as its initial seed.

Second comes preservation of the surrounding stream. Original forward already advanced s0 to s1. After replay, surrounding computation should continue from the state it would otherwise have reached, rather than advance again because of an extra recomputation draw. The diagram and code have no intervening external draws, so that state is s1. Generally, save the surrounding state at replay entry and restore it afterward; do not erase legitimate draws made in between. Restoring s0 alone is incomplete.

The Python random documentation provides getstate and setstate for capturing and restoring a generator’s internal state. Independent Random objects demonstrate the principle below without accessing any real training state.

Nine checks reproduce the replay contract

The code calculates derivatives analytically through the chain rule and uses a fixed local Python seed. bad_rng continues from the state after the original forward. Correct replay begins from the pre-forward state, then restores the surrounding generator to after.

# Scalar chain-rule and Python RNG model; not PyTorch autograd.
from random import Random

def multiplier(rng):
    return 2.0 if rng.random() < 0.5 else 0.0

r = Random(1)
before = r.getstate()
m = multiplier(r)
after = r.getstate()
w, x, target = 2.0, 3.0, 1.0
y = w*x*m
loss = 0.5*(y-target)**2
upstream = y-target
grad = upstream*x*m
bad_rng = Random()
bad_rng.setstate(after)
bad_m = multiplier(bad_rng)
bad_grad = upstream*x*bad_m
assert m == 2.0
assert bad_m == 0.0
assert loss == 60.5
assert grad == 66.0
assert bad_grad == 0.0
r.setstate(before)
replay_m = multiplier(r)
assert replay_m == m
assert upstream*x*replay_m == grad
r.setstate(after)
assert r.getstate() == after
reference = Random()
reference.setstate(after)
assert r.random() == reference.random()
print("9 replay checks passed")

Nine local checks cover the original and alternate multipliers, loss, correct and incorrect gradients, replayed multiplier and derivative, restored state and subsequent random stream. No autograd, real checkpoint, dropout operator or memory benchmark ran. This Python stream must not be confused with arbitrary external random sources that PyTorch automatically manages.

RNG matching is not complete equivalence

The PyTorch reference describes default RNG handling with device-coverage limits. It does not establish the same protection for every custom generator, external library or newly introduced device inside the function. Check the supported scope in the actual version rather than treating RNG preservation as a universal state snapshot.

Other mutable state can invalidate a replay even when its mask matches. A function may read a changing global multiplier, increment a counter that selects a branch or take a different path through an external cache. Restoring random state alone does not rewind those values. A reproducible draw is only one part of a reproducible computation.

Make the region’s dependencies explicit: which inputs determine the result, which state is read-only and which operations have side effects. Logs, counters and cache writes may happen again during recomputation even if tensor values remain correct. Business side effects should rely on neither exactly one invocation nor two complete executions of the region.

Compare more than the forward loss

Start with identical small inputs, parameters and initial random state, then compare forward outputs and parameter gradients with and without checkpointing. Loss agreement alone is insufficient: in this counterexample, that loss was already fixed before the incorrect local derivative was reconstructed. The discrepancy appears later in backward.

Compare gradients parameter by parameter rather than checking only a total norm. Different vectors can share a norm, and errors can cancel during aggregation. Record maximum differences, tolerances, parameter scope and datatype so that close enough has an inspectable meaning rather than being a visual judgment about two training curves.

Add mixed precision, distributed wrappers and real stochastic operators separately when restoring realistic conditions. No measured conclusion about those combinations is supplied here. The earlier AMP clipping article concerns gradient units and operation order; this one asks which computation the derivative belongs to. Both requirements need their own validation.

Measure retained memory and extra work together

Checkpointing reduces retention of selected forward intermediates. It does not imply that parameters, optimizer state, workspaces and every cache shrink by the same ratio. If peak memory comes from another source, checkpointing may leave the limiting allocation largely unchanged.

Hold model, input shapes, batch size, datatype and correctness conditions fixed when recording peak memory and step time. A configuration completing a step establishes execution under that setup. It does not establish a universal saving, nor attribute every throughput change to recomputation. Resource measurements need the same care about conditions as gradient comparisons.

Checkpoint boundaries also matter. Larger regions change the combination of retained values and repeated computation. Locate the source of memory pressure, choose boundaries deliberately and check gradient equivalence before wrapping every layer mechanically. A decision grounded in the actual bottleneck is more useful than a blanket setting copied from another model.

Begin with a region you can explain

A useful diagnostic sequence fixes inputs and parameters, checks a region without randomness or side effects, adds one stochastic operation and validates its replay, then restores the real module and wrappers. When the difference first appears, the newly added condition narrows the investigation.

Do not permanently remove all randomness just to make a comparison pass. Equally, do not classify every tolerance-level difference between implementations as a correctness failure. The aim is to preserve the intended objective and gradients when checkpointing is enabled, rather than force two logs to look identical.

If the gradient counterexample and state timeline help locate a recomputation mismatch, voluntary support through the section below is welcome.

References and catch-up note

Checked October 18, 2026 against the PyTorch checkpoint reference, marked updated July 27, 2026, and Python RNG-state interfaces. This explains an engineering mechanism rather than announcing a release. Both figures, the scalar derivation and nine checks are original teaching material. The pair retains its planned October 18, 09:00 Beijing article date and was posted later that day; framework-specific performance and numerical bounds were not tested.