Activation checkpointing makes GPU memory a scheduling decision

A model that fits for inference may fail before the first training step. The weights did not suddenly grow. Training introduced a second memory problem: autograd must preserve enough of the forward computation to differentiate the loss.

That distinction is easy to lose in advice like “turn on gradient checkpointing.” The flag may be named that way, but the memory being checkpointed is made of activations—intermediate tensors produced by layers.

VRAM is a ledger, not one number

Full fine-tuning commonly allocates parameters, gradients, optimizer moments, possibly higher-precision master weights, saved activations and temporary kernel workspaces. Which category dominates changes with the workload.

Longer sequences and larger micro-batches multiply activation pressure. Adam state scales with the number of trainable parameters. Quantization and LoRA can shrink the model-state side without removing the activation side. Sharding can divide model state across GPUs while each rank still faces local activations.

This is why a memory optimization should name its target:

Method	Primarily changes
Activation checkpointing	saved forward activations
Gradient accumulation	activations per micro-step, while preserving effective batch
LoRA / QLoRA	trainable parameter, gradient and optimizer state
FSDP / ZeRO	model state held by each GPU
Mixed precision	tensor formats and kernel behavior

Turning on every flag is not diagnosis.

Remember boundaries, replay interiors

Ordinary autograd keeps intermediate tensors alive until their backward functions consume them. Activation checkpointing saves selected region inputs and discards tensors inside the region. During backward, the framework runs the region forward again to reconstruct what it needs.

normal:       forward → save every region → backward reads saved tensors
checkpointed: forward → save boundaries   → backward replays each region

The model architecture and mathematical objective are unchanged. The schedule changes: selected forward operations happen twice.

Checkpoint placement controls the trade. Many small regions retain more boundaries but replay less work. A few large regions retain less activation state but replay more. PyTorch’s selective policies can go further by saving expensive operations and recomputing cheaper ones.

Correctness lives inside the replay

The checkpointed function may execute again later. Side effects, changing global state or data-dependent divergence between the original and replayed forward can produce errors or incorrect gradients.

Random operations deserve attention. Dropout should normally regenerate the same mask during recomputation. PyTorch stashes and restores RNG state by default to preserve equivalence, which itself costs time. Moving tensors to a new, unanticipated device inside the region can defeat that guarantee.

Modern PyTorch also asks callers to choose reentrant behavior explicitly. Non-reentrant checkpointing records the graph during forward, supports more autograd patterns and can stop replay once required tensors have been rebuilt. The right call is not to memorize an old snippet; it is to follow the installed PyTorch version’s API and test gradients.

A practical decision sequence

Measure peak allocated and reserved memory on a representative batch.
Estimate model-state memory separately from activation growth across sequence length and micro-batch size.
If activations dominate, checkpoint repeated blocks and measure again.
If model state dominates, use LoRA, sharding, an efficient optimizer or offload instead.
Compare tokens per second, not only seconds per step; the saved memory may permit a larger micro-batch that improves hardware utilization.
Verify loss curves and deterministic tests after changing boundaries.

The deeper idea is useful beyond this one feature: GPU memory is partly a scheduling decision. A tensor can be stored, recomputed, sharded, compressed or offloaded. Each choice moves cost between memory capacity, bandwidth, communication and arithmetic.

Learn the complete mechanism

The interactive Activation Checkpointing lesson shows the forward/backward replay, a relative memory planner, PyTorch and Transformers configuration, and the distinction from accumulation, LoRA and FSDP.