A neural network can fit comfortably in accelerator memory for inference and still run out of memory during training. The reason is that training needs more than model weights. Backpropagation also needs intermediate values from the forward pass, and those activations can consume a large share of memory in deep models or with long sequences and large batches.
Activation checkpointing reduces that memory pressure by deliberately not keeping every intermediate activation. Instead, training saves selected checkpoints and recomputes missing forward-pass values when the backward pass needs them. The trade is straightforward: keep fewer activations in memory, but perform extra computation.
This article builds a practical mental model for that trade. You will learn what checkpointing saves, why recomputation is valid, how checkpoint boundaries affect memory and runtime, and when a different memory optimization is the better first move.
Start with what backpropagation needs
Consider a simple chain of four blocks:
x -> A -> a -> B -> b -> C -> c -> D -> yHere, A through D are parameterized operations and a, b, and c are intermediate activations.
During the forward pass, the network computes y. Training then computes a loss and propagates gradients backward. To calculate gradients for a block, the backward computation often needs values produced during the forward pass. A training framework therefore keeps the required activations alive until their corresponding backward work is complete.
A simplified execution looks like this:
forward: compute a, b, c, y
keep needed intermediates
backward: use c to differentiate D
use b to differentiate C
use a to differentiate B
continue toward AThe exact tensors saved depend on the operations and automatic-differentiation implementation. The important point is that ordinary training preserves forward-pass information because recomputing it is normally avoided.
For a deep network, many such tensors can coexist. Increasing batch size, sequence length, image resolution, hidden width, or layer count can make activation memory large even when parameter memory has not changed much.
The mental model: remember landmarks, rebuild the path
Activation checkpointing changes which forward-pass values are retained.
Suppose blocks B and C form one checkpointed region. Instead of preserving every activation inside that region, training can retain the region input and discard internal values after the forward pass:
forward:
x -> A -> a -> [ B -> b -> C -> c ] -> D -> y
keep a discard b,c as allowedLater, when backward reaches that region, the system runs the necessary part of the forward computation again:
backward reaches checkpointed region
saved a -> B -> b -> C -> c
|
+-> use recomputed values for gradientsThe recomputed activations exist only when needed for that part of backward, rather than being retained from the original forward pass.
This creates the core exchange:
less saved activation state <-> more forward computationCheckpointing does not make the model smaller, reduce the number of trainable parameters, or inherently change the mathematical objective. It changes the execution strategy used to obtain the values required by backpropagation.
See the trade with a small example
Imagine a sequence of six equally sized blocks. For teaching purposes, assume each block produces one activation tensor of the same size and backward needs each one. Real networks are less uniform, but the simplification makes the trade visible.
Without checkpointing, the forward pass keeps six activation tensors:
A1 -> A2 -> A3 -> A4 -> A5 -> A6
keep keep keep keep keep keepNow divide the sequence into two checkpointed regions of three blocks each. Conceptually, training can retain the inputs needed to replay those regions rather than keeping every internal activation for the whole forward-to-backward interval:
checkpoint 1: A1 -> A2 -> A3
checkpoint 2: A4 -> A5 -> A6During backward, each region is recomputed before its gradients are calculated. Peak activation memory can fall because fewer long-lived intermediates span the entire forward pass.
The saving is not simply “six tensors become two tensors.” A real automatic-differentiation system may need to retain outputs, inputs, metadata, non-checkpointed tensors, or other state. Different operations also produce activations of very different sizes. Treat the example as a mental model, not a memory-accounting formula.
The extra work is real too. A checkpointed operation may execute during the original forward pass and again during recomputation. Backward still has to run as usual. Checkpointing therefore reduces memory by spending additional compute time.
Checkpoint boundaries control the trade
A checkpoint boundary determines where retained state ends and recomputation begins. Boundary placement matters because it controls both how much intermediate state can be discarded and how much work must later be repeated.
Consider twelve transformer blocks. One possible policy checkpoints every block. Another groups four blocks into each checkpointed region. A third checkpoints only the memory-heavy middle portion.
These policies can have different behavior even though they train the same architecture.
Finer boundaries
Checkpointing smaller regions can limit how much computation must be replayed at once and gives more control over which operations are checkpointed. But each boundary may require retained tensors and framework bookkeeping. Extremely fine partitioning can therefore provide less benefit than a simple count of checkpoints suggests.
Larger regions
A larger checkpointed region can discard more internal intermediates across that region, but backward may need to replay more computation to reconstruct them. If the region contains expensive operations whose activations are relatively small, the runtime cost may not justify the memory saved.
Uneven networks
Equal numbers of layers do not imply equal memory cost. An attention block handling a long sequence, a high-resolution convolutional stage, and a small projection layer can have very different activation footprints.
For that reason, useful checkpoint placement is usually based on measured memory and compute behavior rather than simply checkpointing every N layers.
Why recomputation must reproduce the required forward values
Checkpointing relies on a basic requirement: the recomputed region must produce the values that backward expects for the training computation being differentiated.
That is easy for deterministic operations whose inputs and parameters have not changed between the original forward execution and recomputation. Stateful or random behavior needs more care.
Suppose a checkpointed region includes dropout. The original forward pass samples a random mask. If recomputation uses an unrelated mask, the reconstructed activations describe a different forward computation. Checkpointing implementations therefore need an appropriate strategy for random-number-generator state when exact replay is required.
Similarly, code whose behavior depends on mutable global state, counters, external side effects, or data that changes between the two executions can make recomputation inconsistent. Framework-specific checkpointing utilities differ in how they preserve random state and what restrictions they impose, so those details should be checked in the framework documentation rather than assumed.
A useful engineering rule is to make checkpointed regions behave like repeatable functions of their explicit inputs, parameters, and controlled randomness.
Estimate whether activations are actually the problem
Checkpointing is useful only when activation memory is large enough to matter. Training memory commonly includes several categories:
parameters
+ gradients
+ optimizer state
+ saved activations
+ temporary workspaces and allocator overheadThe proportions vary with model architecture, numerical precision, optimizer, batch shape, distributed strategy, and framework implementation.
If optimizer state dominates memory, aggressively checkpointing activations may add substantial runtime while freeing too little memory to solve the problem. If activations dominate because sequence length or batch size is large, checkpointing can be much more attractive.
Measure peak memory before choosing the remedy. When possible, inspect memory by training phase or use a framework profiler to identify which tensors and operations remain live. The goal is not merely to reduce a memory number; it is to remove the actual bottleneck with an acceptable performance cost.
Apply checkpointing where memory saved justifies replay
A practical workflow starts with the smallest intervention that makes the workload fit.
First, establish a baseline using the intended batch shape and precision. Record peak device memory and a throughput measure such as training steps or examples per second. A single successful step is not enough if later batches have longer sequences or different shapes.
Next, identify natural repeated regions. Transformer blocks, residual stages, or repeated encoder layers are often convenient boundaries because they have clear inputs and outputs. Convenience alone is not a reason to checkpoint all of them.
Then checkpoint a limited set of memory-heavy regions and measure again. Compare at least:
peak memory
step time or throughput
maximum usable batch or sequence size
training loss behaviorIf the saved memory enables a required sequence length or prevents out-of-memory failures at a modest runtime cost, the trade may be worthwhile. If memory barely changes while step time grows noticeably, checkpoint placement or the optimization itself is probably wrong for that workload.
Finally, test representative training behavior. Dynamic sequence lengths, variable image sizes, different micro-batches, and distributed execution can move the peak to a different point than a small benchmark suggests.
Do not confuse checkpointing with gradient accumulation
Activation checkpointing and gradient accumulation are both used when a desired training workload does not fit in memory, but they solve different parts of the problem.
Gradient accumulation splits a logical batch into smaller micro-batches. Each micro-batch runs forward and backward, and the optimizer updates parameters after gradients from several micro-batches have been accumulated. This can reduce activation memory associated with processing a large batch at once.
Activation checkpointing changes what is retained inside a forward/backward execution. It can help even when a single micro-batch is too large because the model, sequence, or resolution produces too many saved activations.
They can also be combined. For example, gradient accumulation may reduce the micro-batch from eight examples to two, while checkpointing makes those two long-sequence examples fit. The combined runtime cost can be significant, so measure the full training loop rather than assuming that stacking memory techniques is free.
Understand the main failure modes
The most common checkpointing mistakes come from treating memory reduction as the only objective.
Checkpointing cheap activations but expensive computation
A region may take substantial compute while retaining little activation memory. Replaying it buys little memory for a large runtime penalty. Profile before expanding checkpoint coverage.
Ignoring randomness or mutable state
A recomputed region that behaves differently from its original forward execution can invalidate the intended gradient computation. Be especially careful with random operations, stateful modules, and side effects, and follow the guarantees of the checkpointing implementation you use.
Measuring only allocated parameter memory
Checkpointing does not primarily target parameter storage. If the memory report you inspect excludes saved activations or peak temporary allocations, you can draw the wrong conclusion about both the bottleneck and the improvement.
Expecting memory savings to be proportional to checkpoint count
Checkpoint boundaries themselves require retained state, and activation sizes differ across operations. Doubling the number of checkpointed regions does not imply a predictable percentage reduction in peak memory.
Spending the savings without preserving headroom
If checkpointing reduces a run from 95% to 80% of device memory and you immediately increase the batch until usage returns to 99%, small changes in shape or temporary workspace can still cause out-of-memory failures. Production training usually benefits from some memory headroom.
Know when another technique is simpler
Activation checkpointing is a strong option when saved activations are a major memory cost and extra compute is acceptable. It is especially relevant when the model must process long sequences, high-resolution inputs, deep repeated blocks, or a micro-batch that cannot be reduced further without changing the training setup.
It is not automatically the first optimization to apply.
If reducing the micro-batch and using gradient accumulation meets the training objective with acceptable throughput, that can be simpler. If parameter or optimizer state dominates, lower-precision state, optimizer changes, or distributed sharding may target the bottleneck more directly. If temporary kernels or an unexpectedly large tensor cause the peak, changing that operation can be better than replaying broad sections of the network.
Checkpointing also makes less sense when training is already compute-bound and memory has comfortable headroom. Saving memory that the workload does not need only adds recomputation.
The decision should therefore be based on a measured constraint:
Is saved activation memory preventing the required workload from fitting?
yes -> checkpoint selected high-value regions and measure the cost
no -> optimize the resource that is actually limiting trainingConclusion
Activation checkpointing is easiest to understand as a controlled exchange between memory and computation. Ordinary backpropagation keeps forward activations so backward can reuse them. Checkpointing keeps selected landmarks, discards some intermediate state, and reconstructs that state when backward reaches the checkpointed region.
That mechanism is simple, but effective use depends on placement. Measure where training memory goes, checkpoint regions whose saved activations are expensive enough to justify replay, preserve correct behavior for randomness and state, and compare peak memory with actual throughput. When activation memory is the bottleneck, checkpointing can make otherwise impossible training shapes fit without changing the model’s learning objective.