Training a neural network can run out of accelerator memory even when the model parameters fit comfortably. The missing piece is often activations: intermediate values produced during the forward pass and retained because backpropagation needs them later.

Gradient checkpointing, also called activation checkpointing, trades extra computation for lower activation memory. Instead of keeping every intermediate activation until the backward pass, training keeps selected checkpoints and recomputes missing forward values when their gradients are needed.

This article builds a practical mental model for that trade-off, shows what checkpointing changes and what it does not, and explains how to decide whether it is useful for a training workload.

Start with the ordinary training path

Consider a simple network with four sequential blocks:

input -> A -> B -> C -> D -> loss

During inference, an implementation can often discard an intermediate activation after the next operation no longer needs it. Training is different. The backward pass needs information from the forward pass to compute gradients.

A simplified training timeline looks like this:

forward:
input -> A -> B -> C -> D -> loss
          |    |    |    |
          keep keep keep keep

backward:
loss -> grad(D) -> grad(C) -> grad(B) -> grad(A)

Exactly which tensors an automatic-differentiation system saves depends on the operations involved. The important point is that training retains forward-pass state, and that state consumes memory until backward computation can use it.

For deep networks, long sequences, large batches, or large hidden dimensions, saved activations can become a major part of peak training memory.

The checkpointing mental model

Gradient checkpointing changes which forward activations remain stored.

Suppose the network is divided into two regions and only selected boundary values are retained:

input -> A -> B -> C -> D -> loss
          ^         ^
       checkpoint checkpoint

When backward reaches a region whose internal activations were discarded, the training system runs that part of the forward computation again from a retained checkpoint. It then has the intermediate values required to calculate gradients.

Conceptually:

first forward:
compute A, B, C, D
keep selected checkpoints
discard selected intermediate activations

backward through later region:
recompute needed forward values
compute gradients

backward through earlier region:
recompute needed forward values
compute gradients

The benefit is lower activation memory. The cost is additional forward computation during backward.

Checkpointing therefore does not make the same training step cheaper in every dimension. It exchanges one constrained resource for another:

less activation memory <-> more computation

That distinction is the key to using it well.

What checkpointing does not remove

It is easy to interpret “lower training memory” too broadly. Checkpointing primarily targets memory associated with saved activations in checkpointed regions.

A training process may also hold:

  • model parameters;
  • parameter gradients;
  • optimizer state;
  • temporary workspace used by kernels;
  • communication buffers in distributed training;
  • input and output tensors.

Checkpointing does not automatically eliminate these costs. If optimizer state or model parameters dominate memory, activation checkpointing may provide less relief than expected.

This is why the useful question is not simply “Does checkpointing save memory?” It is “How much of this workload’s peak memory comes from activations that can be recomputed?”

A small memory example

Assume a teaching example in which a training step has this approximate device-memory footprint:

parameters, gradients, optimizer state: 12 GB
saved activations:                       10 GB
other temporary memory:                  2 GB
---------------------------------------------
approximate total:                      24 GB

Suppose checkpointing reduces retained activation memory from 10 GB to 4 GB for this particular model and partitioning strategy. The approximate footprint becomes:

parameters, gradients, optimizer state: 12 GB
saved activations:                        4 GB
other temporary memory:                  2 GB
---------------------------------------------
approximate total:                      18 GB

This is deliberately simplified. Real peak memory is not necessarily the sum of a few fixed buckets: tensor lifetimes overlap, allocators reserve memory, temporary workspaces vary, and recomputation itself uses memory.

The example demonstrates the decision principle. A large percentage reduction in activation storage does not imply the same percentage reduction in total training memory.

Checkpoint boundaries control the trade-off

Checkpointing is not only an on-or-off feature. The placement and granularity of checkpointed regions affect both memory and recomputation.

If many activations are retained, backward needs less recomputation but memory savings are smaller. If fewer activations are retained, memory can fall further, but more forward work must be repeated.

For a repeated architecture, such as a stack of transformer blocks, a practical implementation may checkpoint blocks or groups of blocks rather than individual primitive operations. The exact strategy depends on the framework and model architecture.

The right granularity is workload-specific. Measure peak memory and step time rather than assuming the most aggressive checkpointing configuration is preferable.

Why the compute cost is not a fixed percentage

Checkpointed regions execute forward operations again during backward, so training performs additional work. However, there is no universal slowdown percentage.

The observed effect depends on factors such as:

  • how much of the model is checkpointed;
  • the cost of recomputed operations;
  • accelerator utilization;
  • memory bandwidth and kernel behavior;
  • distributed communication;
  • whether lower memory enables a different batch size or sequence length.

A configuration that adds recomputation can still improve the usefulness of a fixed-memory device if it enables the batch size, sequence length, or model configuration the workload requires. Conversely, checkpointing is unnecessary overhead when the training job already fits with comfortable memory headroom.

Evaluate both peak memory and time per training step under the same workload.

Recomputed forward passes must preserve training semantics

Recomputation introduces an important correctness requirement: the repeated forward computation must behave consistently with what backward expects.

This matters when a checkpointed region contains stateful or stochastic behavior. For example, dropout uses random masks during training. If recomputation used unrelated randomness without the framework accounting for it, the repeated forward path could differ from the original one in a way that changes gradient computation.

Checkpointing implementations can preserve or restore random-number-generator state to handle such cases, but the guarantees and controls are framework-specific. Do not assume every custom checkpointing mechanism handles randomness identically.

State mutation is another warning sign. A forward function that changes external state, advances counters with semantic meaning, performs side effects, or depends on mutable values may behave differently when executed more than once.

A useful design rule is:

checkpointed forward computation should be safe to recompute

When using a framework’s checkpointing API, read its documented restrictions around randomness, mutable state, in-place operations, and automatic differentiation rather than treating recomputation as invisible.

Checkpointing does not change the optimization goal

Activation checkpointing is a memory-management technique, not a new learning objective. In a correctly configured deterministic case, recomputing an intermediate value should provide the same value needed for the same gradient calculation, apart from normal numerical effects of the execution environment.

That means checkpointing is conceptually different from techniques such as reducing batch size, truncating a sequence, shrinking the model, or changing numerical precision. Those changes can alter the training workload itself or its numerical behavior. Checkpointing instead aims to retain the training computation while changing when selected intermediate results are produced and stored.

Implementation details still matter. A bug involving state, randomness, or unsupported operations can violate that intent, which is why correctness checks belong in the rollout process.

Use measurement to decide whether checkpointing helps

A practical workflow starts by identifying the actual memory bottleneck.

First, run the target model, batch size, sequence length, precision, and optimizer without checkpointing if the configuration fits at all. Record peak device memory and step time. If it does not fit, use a smaller configuration to establish a baseline and inspect memory breakdowns available from the framework or profiler.

Next, enable checkpointing for a sensible model region and repeat the same measurement. Compare:

peak allocated/reserved memory
step time
examples or tokens processed per second
final feasible batch or sequence configuration

If checkpointing makes a previously impossible target configuration fit, the extra compute may be an acceptable price. If memory barely changes, investigate whether activations are actually the dominant consumer or whether the chosen regions retain more state than expected.

Finally, run a short correctness comparison. Check that losses and gradient behavior remain plausible relative to the non-checkpointed setup under controlled seeds and equivalent settings. Bit-for-bit equality is not a general requirement across all hardware and kernels, but large unexplained divergence deserves investigation.

Common mistakes

Treating checkpointing as free memory

Memory savings come from recomputation. A training plan that ignores the extra compute can underestimate wall-clock cost.

Expecting it to solve every out-of-memory error

If parameters, optimizer state, or a temporary operation dominates peak memory, checkpointing may not address the limiting allocation. Other techniques target different memory components.

Changing several variables at once

Enabling checkpointing while also changing batch size, precision, and sequence length makes the result difficult to interpret. Establish a comparable baseline, then change one major factor at a time when possible.

Ignoring stochastic or stateful operations

A checkpointed region can execute its forward logic more than once per training step. Verify the framework’s behavior for random state and avoid unintended side effects.

Measuring only whether the model fits

A configuration that fits but becomes impractically slow may not be useful. Memory, throughput, and model-quality checks all belong in the evaluation.

When to use it

Gradient checkpointing is a strong candidate when saved activations are a meaningful part of peak memory and memory prevents you from using the required model, sequence length, or batch configuration. It is especially relevant for deep networks where many intermediate activations would otherwise remain live until backward.

It is less compelling when training already fits comfortably, recomputation is too expensive for the available compute budget, or non-activation memory is the real bottleneck. In those cases, checkpointing adds complexity and compute without addressing the main constraint.

Conclusion

Gradient checkpointing is easiest to understand as a scheduling decision for intermediate results. Ordinary training stores activations so backward can reuse them. Checkpointed training stores fewer of them and regenerates selected values later.

That gives developers a practical lever when activation memory is scarce: spend additional computation to reduce retained forward state. The useful configuration is not the one that checkpoints the most operations. It is the one that meets the workload’s memory requirement while keeping training time and correctness within acceptable bounds.

Artificial Intelligence Neural Networks Model Training Gradient Checkpointing