Activation Checkpointing — Trading Compute for GPU Memory

A 7-billion-parameter model can fit its quantized weights (stored in fewer bits per number) on a GPU and still run out of memory the moment training begins. Inference (using the model, not training it) mainly needs weights and a working set such as the KV cache (stored attention keys and values for text already read). Training must also keep the information needed to differentiate the loss (to find how each weight should change).

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. It usually grows with micro-batch size × sequence length × hidden width × depth: how many examples run at once, how long each is, how wide each layer is, and how many layers there are.

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 constant factors change with:

• attention implementations
• fused kernels
• tensor parallelism
• padding
• allocator behavior
• framework versions

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.

It 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. Actual overhead depends on:

• how much is checkpointed
• the forward/backward ratio
• kernel efficiency
• communication overlap
• 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. Gradients may be wrong or PyTorch may raise an error if it:

• mutates global state
• reads changing external state
• follows a different branch during recomputation

Dropout (randomly zeroing some values during training) needs special treatment because recomputation should reproduce the same mask as the original forward pass. PyTorch preserves and restores the random-number generator (RNG) 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.

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.