What is activation checkpointing and how does it save GPU memory?

Activation checkpointing reduces GPU memory used to save forward activations by keeping selected boundary tensors and recomputing omitted intermediate activations during backpropagation. It trades additional computation for less activation storage, often enabling a larger batch, longer sequence, or deeper model to fit in GPU memory. The method is also called *gradient checkpointing*, although it primarily targets saved forward intermediates—not the gradient or optimizer-state buffers.

1. Why does an ordinary neural network retain activations?

• A network processes x₀ → L₁ → x₁ → L₂ → x₂ → ... → L₁₆ → x₁₆ during its forward pass.
• Backpropagation needs information from earlier computations, so autograd normally saves relevant layer inputs, outputs, masks, statistics, and other tensors.
• The amount saved depends on the operation. It is not literally one activation per layer in every real model.
• Long Transformer sequences, high-resolution CNN features, large batches and deep networks can make activations consume substantial GPU memory.
• Activation checkpointing deliberately does not retain some of these tensors; it obtains them again when backward reaches the corresponding region.

Think of it as keeping bookmarks in a long calculation instead of saving every intermediate page. When backward needs a missing page, the system starts from an earlier bookmark and recalculates it.

2. How does checkpointing work during forward and backward?

1. Partition the network: Choose groups of layers or operations called checkpointed segments.
2. First forward pass: Execute the model while retaining selected segment inputs/outputs and required metadata, instead of all saved internal tensors.
3. Backward visits a checkpointed region: It reruns enough forward operations from the retained input to rebuild needed tensors.
4. Calculate gradients: Autograd uses the regenerated intermediates for the usual backward differentiation.
5. Release temporary work: The segment's recomputed values can be discarded as backward proceeds to earlier segments.

Illustrative 16-layer layout:

• Forward: x₀ → [L₁…L₄] → x₄ → [L₅…L₈] → x₈ → [L₉…L₁₂] → x₁₂ → [L₁₃…L₁₆] → x₁₆.
• Retained checkpoint boundaries in the simple model: `x₀, x₄, x₈, x₁₂, x₁₆`.
• Backward in the final segment reconstructs x₁₃, x₁₄, x₁₅ and any other intermediates its derivative computations need from x₁₂; then proceeds toward the preceding segments.

An activation checkpoint is not a checkpoint file that saves model weights for later restoration.

3. Beginner-friendly numerical example: A 16-layer network

Assume 16 identical layers and one activation-sized tensor of 8 MiB at each layer boundary, including input x₀. Here 1 MiB = 1,048,576 bytes. For this example only, ignore model parameters, gradients, optimizer states, additional operator-saved tensors, workspaces and allocator overhead.

Step A — Without checkpointing

• There are 17 boundary tensors: the initial input plus 16 outputs.
• Simplified retained-activation storage: 17 × 8 MiB = 136 MiB.

Step B — Split into four segments of four layers

• Save only five boundaries: x₀, x₄, x₈, x₁₂, x₁₆.
• Saved-boundary storage: 5 × 8 MiB = 40 MiB.
• Difference immediately after initial forward: 136 − 40 = 96 MiB of illustrative retained-tensor storage.

Step C — Allow memory for recomputation

While backward processes one segment, suppose at most four additional activation-sized tensors must temporarily coexist.

• Temporary scratch: 4 × 8 MiB = 32 MiB.
• Checkpoint boundaries plus scratch: 40 + 32 = 72 MiB.
• Illustrative activation-related peak: 72 MiB, rather than 136 MiB.

Step D — Calculate the illustrative reduction

• Activation budget difference: 136 − 72 = 64 MiB.
• Percent reduction: 64 / 136 × 100% ≈ 47.1%.

Critical limitation: The 47.1% is not a guaranteed overall GPU-memory saving. The calculation assumes at most four simultaneously live scratch tensors, and counts five retained boundaries for the whole backward pass. Actual PyTorch peak memory depends on the saved tensors each operator requires, when boundaries are freed, any recomputation early stopping, kernel temporaries, parameters, optimizer state, batch shape and allocator behavior. Only a profiler can establish real peak memory.

4. What extra computation does this require?

Checkpointing still performs the initial forward and normal backward work, but it may repeat part of the forward work.

Toy arithmetic:

• No checkpointing: 16 original layer-forward evaluations.
• Full segment replay: up to approximately 16 *additional* layer-forward evaluations.
• This can mean 100% extra forward-layer work, but not 100% longer training because backward and other work already consume time.

For a separate *illustrative cost model*, assume a layer's forward costs 1 compute unit and backward costs 2.

• Baseline: 16 × (1 + 2) = 48 units.
• With one extra forward replay per layer: 16 × (1 + 1 + 2) = 64 units.
• Extra total compute: (64 − 48) / 48 × 100% ≈ 33.3%.

This 33.3% is not a measured time overhead. PyTorch's non-reentrant checkpoint implementation can stop replaying a region once required tensors are recovered (*early stop*), so it may execute less than a complete segment replay; kernels, bandwidth and RNG-state bookkeeping further change wall-clock results.

5. Choosing the segment size

For a simplified linear network of N equally costly layers, checkpoint groups of k layers have:

• Approximately N/k saved segment boundaries.
• Approximately k activation units needed for reconstruction scratch.
• Idealized activation storage scaling like `O(N/k + k)`.
• Choosing k ≈ sqrt(N) balances the two terms, giving `O(sqrt(N))` activation storage in this simplified setup, rather than `O(N)`.

But practical layers vary greatly in size and cost. A residual network, attention block or branched graph does not necessarily benefit from equal numbers of layers per segment.

Choose boundaries by saved activation bytes versus recomputation cost, not by layer count alone. Start with a few activation-heavy blocks and benchmark several segment sizes.

6. Why checkpoint Transformer blocks?

• Transformers retain attention, feedforward and normalization-related intermediates over many layers and tokens.
• Longer context length and larger batches often increase saved activation usage.
• Checkpoint whole Transformer blocks, several blocks together, or selected high-memory sub-operations.
• Check whether FlashAttention or fused kernels already avoid storing large intermediate attention matrices; the incremental checkpoint benefit changes accordingly.
• Fine-tuning with trainable adapters can still require saved inputs from substantial portions of the forward computation.
• Extra reconstruction compute can reduce tokens per second; checkpoint only what is needed to meet the memory target.

7. Modern PyTorch implementation

Current PyTorch documentation recommends torch.utils.checkpoint.checkpoint(..., use_reentrant=False). This example checkpoints groups of four blocks, while protecting the group reference against late-bound Python closure errors when backward reruns the code.

PYTHON
import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint

class CheckpointedStack(nn.Module):
    def __init__(self, blocks, group_size=4):
        super().__init__()
        self.blocks = nn.ModuleList(blocks)
        self.group_size = group_size

    def forward(self, x):
        for start in range(0, len(self.blocks), self.group_size):
            group = tuple(self.blocks[start:start + self.group_size])

            def run_group(inp, layers=group):
                for layer in layers:
                    inp = layer(inp)
                return inp

            x = checkpoint(
                run_group,
                x,
                use_reentrant=False,
                preserve_rng_state=True
            )
        return x

# Example: blocks is a list of compatible neural-network modules.
model = CheckpointedStack(blocks, group_size=4).cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

optimizer.zero_grad(set_to_none=True)
out = model(inputs.cuda())
loss = loss_fn(out, labels.cuda())
loss.backward()
optimizer.step()

Important details:

• `use_reentrant=False`: Recommended non-reentrant variant. It supports more autograd workflows, including torch.autograd.grad, and normally uses early stopping during recomputation.
• `layers=group`: Binds the correct blocks at each loop iteration, preventing an incorrect last-group replay during backward.
• `preserve_rng_state=True`: Keeps relevant random behavior consistent across initial forward and replay.
• No weight-copy trick: This procedure does not change parameter or optimizer-state bytes.
• Alternative helper: checkpoint_sequential can split sequential models into segments, but its documented behavior generally does not checkpoint the final segment. Verify which portions are actually checkpointed.
• The example assumes compatible modules, inputs, labels, loss and CUDA hardware are available; it is a training pattern, not a runnable pretrained model download.

8. RNG-state preservation and dropout

Dropout samples a random mask. If recomputation uses a different mask from the initial forward, regenerated activations may differ, producing different or incorrect gradients.

• PyTorch's default preserve_rng_state=True saves/restores relevant RNG state for consistency with a non-checkpointed execution.
• preserve_rng_state=False can save RNG bookkeeping, but the stochastic recomputation may then use different draws and change the computed gradients. Use only if that behavior is explicitly acceptable.
• PyTorch documents CPU state plus one additional inferred device type; moving computations to an otherwise untracked device inside a checkpointed function can undermine RNG equivalence.
• Under torch.compile, the documentation notes that preserve_rng_state=False is ineffective: RNG state remains preserved.
• Non-reentrant determinism_check can check properties such as shape, dtype and device; it does not prove bitwise equality of every recomputed value.
• Avoid changing training mode, module weights, global flags, mutable buffers, or side-effect behavior between forward and recomputation.

Validate checkpointed gradients against ordinary backpropagation on a small reproducible test before scaling up.

9. Interaction with AMP and loss scaling

Activation checkpointing and automatic mixed precision address different dimensions of the memory problem:

• AMP: Uses lower-precision arithmetic for suitable operations and may reduce some tensor sizes, as in Questions 293–295.
• Checkpointing: Saves fewer activations and recomputes what backward needs.
• Together: Both may reduce activation memory, but savings are not guaranteed to multiply; recomputation, tensor casts and FP32 operations still consume resources.
• Keep autocast behavior consistent: When recomputed, operations should see the same relevant dtypes and computational context as in the original forward; do not change model mode or parameters mid-graph.
• GradScaler ordering stays the same: If using FP16 loss scaling, complete backward, unscale and clip correctly before the optimizer step. Checkpointing does not replace a scaler.

10. Distinguishing related memory optimizations

• Gradient accumulation: Splits a larger effective batch into smaller microbatches and accumulates gradients across them. It lowers per-microbatch activation needs, but does not reconstruct deleted activations within a backward pass.
• Optimizer-state savings: Reducing the precision or sharding the optimizer's moment buffers targets optimizer memory, not forward activations.
• CPU offloading: Moves selected tensors or states off the GPU and back when needed, exchanging GPU residency for CPU memory and transfer overhead. Unlike checkpointing, it need not recompute the forward operations.
• Saving a training checkpoint to disk: Stores weights and optimizer state for later recovery; unrelated to per-step activation checkpointing.

These techniques can be combined, but each has its own memory, computation and correctness costs.

11. How to profile and select a useful checkpoint plan

1. Measure baseline: Identify whether activations, parameters, optimizer state or temporary attention buffers are responsible for the out-of-memory problem.
2. Choose high-value regions: Begin with blocks that save substantial activations but are relatively cheap to recompute.
3. Test coverage and grouping: Try whole-block checkpointing, every second block or selected attention/MLP modules, and compare actual savings.
4. Verify numerical correctness: On fixed inputs, compare losses and gradients against a non-checkpointed baseline within appropriate floating-point tolerances.
5. Preserve randomness: Hold seeds, dropout behavior and training mode consistent; check RNG behavior explicitly when needed.
6. Measure peak GPU memory: Use torch.cuda.reset_peak_memory_stats() and torch.cuda.max_memory_allocated(); optionally inspect max_memory_reserved().
7. Measure throughput and latency: Use warmed-up GPU runs and synchronization. Record step time and tokens/images per second.
8. Control other variables: Keep batch size, input shape, optimizer, precision policy, data and CUDA environment identical for initial comparisons.
9. Retain headroom: Variable sequences or temporary allocations can cause later OOM even if one test step fits.

Tiny CUDA profiling sketch:

PYTHON
torch.cuda.synchronize()
torch.cuda.reset_peak_memory_stats()
# Execute representative warmed-up forward + backward here.
torch.cuda.synchronize()
peak_mib = torch.cuda.max_memory_allocated() / (1024**2)
print("Peak allocated CUDA tensor memory:", peak_mib, "MiB")

The reported allocation excludes some external GPU runtime memory. Peak reserved memory may exceed peak allocated memory due to the caching allocator.

12. Limitations and common mistakes

• Memory claims without profiling: The example's 47.1% activation-only reduction is not a real-model or total-GPU benchmark.
• Checkpointing the wrong region: Cheaply stored activations with costly recomputation can cause a net performance loss.
• Assuming no compute cost: Saved forward work must generally be recomputed during backward.
• Incorrect loop closure: Capturing mutable or late-bound layer references can replay the wrong module and silently corrupt gradients.
• Dropout mismatch: Disabling RNG preservation without considering dropout changes recomputation and may break gradient equivalence.
• Unexpected side effects: Stateful operations or mutations during recomputation can produce inconsistent values.
• Assuming all devices are tracked: RNG handling has device restrictions.
• Overlooking fused kernels: Modern attention kernels can alter both baseline saved-activation memory and recomputation trade-offs.
• Missing autograd requirements: Some legacy reentrant APIs have more constraints. Prefer explicit use_reentrant=False.
• Assuming optimizer states shrink: Checkpointing doesn't reduce the number of model weights or optimizer moments.
• Ignoring throughput: A memory reduction may come with unacceptable iteration-time overhead; compare both memory and end-to-end training.
• Incompatible recompute context: Different AMP dtypes, model configurations or randomness can cause gradient errors.

Curriculum boundary: Questions 293–295 introduced mixed precision, FP16/BF16 and gradient loss scaling. Question 296 focuses on activation-saving/recomputation. Question 297, "What is quantization-aware training (QAT)?", begins the next efficiency topic. Do not conflate activation checkpointing with reduced-precision training or optimizer-state compression.

Key takeaway: Activation checkpointing saves memory by retaining fewer forward intermediates and recalculating missing activations during backward. It can markedly lower activation memory at the cost of extra computation, so profile the saved tensors, numerical correctness and time trade-off together.

Research and technical references: Chen et al., *Training Deep Nets with Sublinear Memory Cost*; PyTorch, *torch.utils.checkpoint*; PyTorch, *Current and New Activation Checkpointing Techniques*.