Activation Checkpointing — Trading Compute for GPU Memory

A 7-billion-parameter model can fit its quantized weights on a GPU and still run out of memory the moment training begins. Inference mainly needs weights and a working set such as the KV cache. Training must also retain information needed to differentiate the loss.

The important first question is therefore not “How large are the weights?” It is “Which category is filling memory?”

The training-memory ledger

Activation memory depends on architecture, dtype and implementation, but it usually grows with micro-batch size × sequence length × hidden width × depth. For long-context fine-tuning, those saved tensors can rival or exceed the model state. This is why enabling checkpointing may rescue an out-of-memory run at 8,192 tokens while barely helping a short-sequence run dominated by Adam state.

The picture above is a relative sketch. Real VRAM must be measured because attention implementations, fused kernels, tensor parallelism, padding, allocator behavior and framework versions change the constant factors.

What normal backward expects

Suppose a four-block network computes:

x0 → block1 → x1 → block2 → x2 → block3 → x3 → block4 → loss

Without checkpointing, autograd saves tensors such as x1, x2 and x3. During backward, each block uses the saved inputs or outputs to calculate local derivatives. Keeping them is fast: backward reads rather than re-running the forward block. It is also memory hungry because early activations must survive until backward finally reaches them.

With block 1 and block 3 selected as checkpoint boundaries, the system keeps the boundary inputs and discards intermediates inside each region. When backward reaches the second region, it replays that region’s forward operations from the saved boundary, reconstructs the tensors, uses them for gradients and releases them again.

The computation graph is not removed. It is evaluated twice in selected regions: once during the original forward pass and once as backward requests the missing values.

The trade-off is adjustable

Checkpoint every block and you save many boundaries but recompute short regions. Checkpoint only a few large regions and you retain fewer boundaries but replay more computation. Selective checkpointing can keep expensive operations while recomputing cheap ones.

PyTorch’s current API also distinguishes reentrant and non-reentrant behavior; the documentation recommends passing use_reentrant explicitly. The non-reentrant implementation can stop recomputation once the needed tensors have been reconstructed and supports more autograd patterns.

There is no universal “20% slower” law. Hugging Face documents roughly 20% as a useful expectation for common Transformer training, but actual overhead depends on how much is checkpointed, the forward/backward ratio, kernel efficiency, communication overlap and whether memory savings permit a better micro-batch. Benchmark tokens per second, not just step time.

PyTorch: checkpoint a pure region

import torch
from torch.utils.checkpoint import checkpoint

class CheckpointedBlock(torch.nn.Module):
    def __init__(self, block):
        super().__init__()
        self.block = block

    def forward(self, hidden_states):
        return checkpoint(
            self.block,
            hidden_states,
            use_reentrant=False,
        )

The checkpointed function may run again during backward. Treat it like a pure function of its inputs. If it mutates global state, reads changing external state, or follows a different branch during recomputation, gradients may be wrong or PyTorch may raise an error.

Dropout needs special treatment because recomputation should reproduce the same mask as the original forward pass. PyTorch preserves and restores random-number generator state by default. That preservation has overhead, but disabling it is only safe when equivalent deterministic behavior is not required. Moving tensors to an unanticipated new device inside a checkpointed function can also break RNG equivalence.

Transformers: one flag, same mechanism

from transformers import TrainingArguments

args = TrainingArguments(
    output_dir="runs/demo",
    per_device_train_batch_size=2,
    gradient_accumulation_steps=8,
    bf16=True,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"use_reentrant": False},
)

For supported models, model.gradient_checkpointing_enable() exposes the same idea directly. A library flag is convenient, but verify it actually covers the expensive model blocks and measure peak allocated/reserved CUDA memory around a representative training step.

Do not confuse five different memory tools

These methods compose. A QLoRA run may use quantized frozen base weights, low-rank trainable adapters, small micro-batches, gradient accumulation and activation checkpointing together. Each line solves a different part of the ledger.

Safe measurement lab

The proxy below models only activation units — saving every k-th boundary exchanges retained tensors for recomputation:

import math

layers = 32
checkpoint_every = 4
activation_mb_per_layer = 180

normal_mb = layers * activation_mb_per_layer
saved_boundaries = math.ceil(layers / checkpoint_every) + 1
checkpointed_mb = saved_boundaries * activation_mb_per_layer
recomputed_layers = layers - math.ceil(layers / checkpoint_every)

print(f"normal saved activations: {normal_mb:,} MB")
print(f"checkpointed proxy:       {checkpointed_mb:,} MB")
print(f"relative activation use:  {checkpointed_mb / normal_mb:.1%}")
print(f"block outputs replayed:   about {recomputed_layers}")
print("Weights, gradients and Adam state are unchanged.")

normal saved activations: 5,760 MB
checkpointed proxy:       1,620 MB
relative activation use:  28.1%
block outputs replayed:   about 24
Weights, gradients and Adam state are unchanged.

Production checklist

• Record peak VRAM and tokens/second with and without checkpointing on the same batch, sequence distribution and warm-up period.
• Checkpoint repeated, memory-heavy blocks; avoid tiny regions whose scheduling and RNG overhead outweigh memory saved.
• Pass use_reentrant deliberately and follow the PyTorch version’s current recommendation.
• Keep checkpointed regions deterministic with respect to their inputs.
• Test dropout and multi-device code; do not casually disable RNG preservation.
• Combine with accumulation only after distinguishing micro-batch size from effective batch size.
• Validate gradients or loss curves after changing checkpoint boundaries.
• Measure reserved and allocated memory; fragmentation can look like activation pressure but needs allocator or shape fixes instead.

In one breath

• Training VRAM holds parameters, gradients, optimizer state, and saved activations — checkpointing only shrinks the last one.
• Activation memory grows with micro-batch × sequence length × hidden width × depth, so it dominates long-context fine-tuning.
• Checkpointing keeps a few boundary activations and discards the intermediates, then recomputes them during backward — trading extra forward compute for memory.
• It composes with mixed precision, gradient accumulation, LoRA/QLoRA, and FSDP — each targets a different line of the memory ledger.
• It can’t help if parameters + gradients + optimizer state alone exceed the GPU; diagnose the category first, keep checkpointed regions pure and deterministic, and benchmark tokens/second.

Quick check

Sources and next step

• PyTorch torch.utils.checkpoint
• PyTorch activation-checkpointing techniques
• Hugging Face gradient checkpointing

Next, learn when model state—not activations—is the bottleneck in Multi-GPU training with DDP and FSDP.