activation checkpointing makes GPU memory a scheduling decision

At 2:17 a.m., a training job dies with CUDA out of memory during the first backward pass. The model ran inference on the same GPU five minutes earlier.

Nothing mysterious happened to the weights. Training added a bill that inference never had to pay: autograd kept intermediate results from the forward pass because backward needs them to calculate gradients.

My strong opinion is this: activation checkpointing is not a “make the model smaller” switch. It is a scheduling decision. You choose to store a tensor now, or spend arithmetic later to recreate it. That makes it powerful when saved activations are the bottleneck, and almost useless when optimizer state or parameters are the bottleneck.

The difference matters. I have watched teams turn on gradient checkpointing, celebrate a lower memory number, and then discover that their optimizer states still consume nearly the entire GPU. The flag was working. The diagnosis was not.

VRAM is a ledger, not one number

A training process can hold several different kinds of data:

• Parameters are the model’s learned weights.
• Gradients are the derivatives calculated for those weights.
• Optimizer state is extra memory such as Adam’s first and second moments.
• Activations are intermediate tensors produced while the model runs forward.
• Workspaces are temporary buffers used by kernels.
• Allocator reservations are GPU blocks held by the framework for reuse.

Inference mainly needs parameters, plus temporary tensors and sometimes a key-value cache for generation. Training needs parameters, gradients, optimizer state, and enough activations to run backward. That is why “it fits for inference” proves very little about training.

Consider a rough full fine-tuning example: a 1.3-billion-parameter Transformer using bfloat16 parameters and gradients with Adam. Assuming Adam stores two FP32 moment tensors, the arithmetic looks like this:

• Parameters: 1.3 billion multiplied by 2 bytes is about 2.6 GB, or 2.42 GiB.
• Gradients: another roughly 2.42 GiB if stored in bfloat16.
• Adam’s two 32-bit moment tensors: about 10.4 GB, or 9.69 GiB.
• A 32-bit master copy of the parameters: another 5.2 GB, or 4.84 GiB, if the optimizer keeps one.

Under the FP32-moments-plus-master-copy assumption, that subtotal is about 19.4 GiB before saved activations, temporary workspaces, CUDA context overhead, and fragmentation. On a 24 GiB GPU, there is not much room left for a long sequence.

Those dtypes and the master copy are not universal. Some optimizer implementations store moment tensors in the parameter dtype; others use FP32 or a compressed format, and a master copy may or may not exist. Check the parameter, gradient, moment, and master-copy dtypes for the actual optimizer implementation. The point is not that every 1.3-billion-parameter model consumes exactly 19.4 GiB. The point is that model-state memory can be estimated separately from activation memory, and the two respond to different remedies.

A longer sequence increases token-shaped activations roughly in proportion to the number of tokens. A larger micro-batch does the same. Attention introduces another wrinkle: a naïvely materialized attention matrix grows with the square of sequence length, while memory-efficient attention kernels avoid storing that full matrix. Do not diagnose sequence memory from the words “attention is quadratic” alone. Profile the actual kernels.

LoRA and QLoRA make the same distinction from the other direction. They can dramatically reduce trainable-parameter gradients and optimizer state, but the base model still runs through every token. Its forward activations can remain large. Sharding can divide model state across GPUs while leaving each rank with its own local activation workload.

The useful question is never “Which memory trick should we enable?” It is “Which line in the ledger is crossing the capacity limit?”

What checkpointing actually changes

Ordinary autograd saves intermediate tensors as the forward pass runs. A later backward operation consumes those tensors. For a linear layer, for example, calculating the weight gradient needs the layer’s input. For a nonlinear operation, backward may need the input or the output to evaluate its derivative.

Without checkpointing, a simplified execution looks like this:

forward → save intermediates from every layer → backward reads them

Activation checkpointing marks a region of the forward pass as disposable. The framework keeps the region’s boundary inputs, but releases many internal tensors after the first forward. During backward, it runs that region again and reconstructs the missing intermediates.

forward → save region boundaries → discard interiors
backward → replay region → use reconstructed intermediates

The model’s function has not changed. The order and number of operations have changed. Selected forward operations now happen twice.

That is why “gradient checkpointing” is a slightly misleading name. The gradients are not being checkpointed. The system is checkpointing enough information about the forward computation to reconstruct the activations needed for gradients.

It is also different from saving a model checkpoint to disk. A training checkpoint is a persistence mechanism used for resuming a job. Activation checkpointing is an execution strategy used inside one forward and backward pass. Same word, very different object.

Suppose a Transformer block is wrapped in a checkpoint. The initial forward keeps the block input and returns its output, but does not retain every internal projection, gate, and normalization tensor. When backward reaches that block, the framework re-enters the block with gradients enabled. It stops replay when it has reconstructed the values required by the pending backward operations.

The boundary is the important part. Checkpointing a whole stack of blocks saves fewer boundary tensors than checkpointing every tiny operation, but replaying a large region creates a larger temporary peak during recomputation. Checkpointing many small regions creates more boundaries and framework overhead. There is no universally best granularity.

Modern PyTorch also supports more selective policies: retain expensive-to-recompute operations and discard cheaper ones. That is often a better schedule than blindly wrapping every operation, but it requires measurement. A memory policy that looks elegant on paper can lose to a simpler block-level policy once kernel overhead and communication enter the picture.

A worked memory example

Return to the 1.3-billion-parameter model. Assume its Transformer blocks have hidden width 2,048, the micro-batch contains 2 sequences, the sequence length is 2,048 tokens, and the activation format uses 2 bytes per value.

One hidden-width tensor has:

2 × 2048 × 2048 × 2 = 16,777,216 bytes

That is exactly 16 MiB using binary units.

Now suppose a profiler attributes six hidden-width-sized tensors to each of 24 blocks. This is deliberately simplified; real blocks also contain wider feed-forward intermediates and kernel-specific buffers.

The retained activation estimate is:

24 × 6 × 16 MiB = 2304 MiB

That is 2.25 GiB. A feed-forward tensor with four times the hidden width is four times larger than one of those hidden-width tensors. Whether it survives depends on the operation and autograd implementation, so this arithmetic is an intuition builder, not a peak-memory prediction.

If each block is checkpointed, the simplified boundary cost is closer to:

24 × 16 MiB = 384 MiB

The block’s internal tensors still exist temporarily while that block is recomputed and differentiated. Checkpointing has not made activations free. It has changed the peak from “many layers’ interiors survive together” to “boundaries survive, and one region’s interiors exist during replay.”

That can be enough to fit on the 24 GiB card. It can also fail if the 19.4 GiB model-state estimate was already optimistic, if a fused kernel allocates a large workspace, or if attention materializes an unexpected score tensor.

This is why a percentage such as “checkpointing costs 20 percent” is not a property of the technique. If profiling shows that a region’s forward takes 40 milliseconds and its backward takes 80 milliseconds, replaying that forward adds 40 milliseconds before other overhead. The region’s 120-millisecond contribution becomes 160 milliseconds, a 33 percent increase. Another model may have a different forward-to-backward ratio, different kernels, or a larger batch after checkpointing.

Measure tokens per second, not just seconds per optimizer step. A checkpointed run that permits a micro-batch of 2 instead of 1 may process more tokens per second despite doing extra arithmetic. Or it may not. GPUs are annoyingly literal about utilization.

Replay is part of your program

The checkpointed function runs more than once. That makes correctness depend on what the function does besides pure tensor computation.

Dropout is the classic example. The initial forward samples a dropout mask. During replay, backward needs the same mask to reproduce the same mathematical path. PyTorch saves and restores random-number-generator state by default for the CPU and the relevant accelerator device type. That preserves the expected behavior, but saving and restoring state costs time.

Moving tensors to a new, unanticipated device inside the checkpointed function can defeat that guarantee. So can a function that reads mutable global state, increments a counter, changes a cache, calls an external service, or chooses a different branch on replay.

A safe checkpointed region is close to a pure function: tensor inputs go in, tensor outputs come out, and the same inputs produce the same computation. A custom layer that updates a running statistic or consumes a one-time iterator deserves suspicion.

For current PyTorch APIs, pass the reentrant choice explicitly rather than relying on an old snippet copied from a forum. Non-reentrant checkpointing is generally the more capable option: it records more autograd behavior, supports more backward patterns, and can stop replay once the required tensors are available. Reentrant checkpointing has stricter input and output requirements. Check the documentation for the PyTorch version installed in your environment.

A minimal block-level pattern is:

from torch.utils.checkpoint import checkpoint

hidden = checkpoint(
    run_transformer_block,
    hidden,
    use_reentrant=False,
)

The exact arguments depend on the block. Attention masks, position information, and other tensor inputs must be passed consistently on the first execution and the replay. If a framework exposes a model-level gradient-checkpointing switch, inspect what regions it wraps instead of assuming it matches your intended boundaries.

For decoder-only language models, disable key-value caching during training. A cache is useful when generating tokens one at a time, but it is not the training activation strategy you want, and many Transformer implementations treat caching as incompatible with gradient checkpointing.

The strongest objection is often right

The best counterargument is simple: use a larger GPU, lower the precision, use a memory-efficient attention kernel, or shard the model. Those options may reduce memory without recomputing the forward pass.

If the bottleneck is optimizer state, checkpointing is the wrong tool. Use an optimizer with a suitable state layout, parameter-efficient fine-tuning, sharding, or offload. If the bottleneck is a materialized attention matrix, fix the attention implementation first. If the job is already well within the memory budget, checkpointing buys nothing and adds work.

Checkpointing earns its place when saved activations dominate after those basics are handled, especially when the alternative is reducing the micro-batch so far that the GPU sits idle. It is also attractive when adding another GPU would introduce communication, scheduling, and operational cost disproportionate to the arithmetic being recomputed.

The trade is not memory versus nothing. It is memory capacity versus compute, random-state handling, kernel overhead, and sometimes distributed communication. Treat it as a budget decision.

What to do on Monday morning

Start with a representative training step: the real sequence length, labels, attention mask, precision, optimizer, and micro-batch. Warm up first because kernels and optimizer state may allocate lazily. Then reset peak statistics and measure several identical steps.

For a Transformers-style model whose batch contains labels, a small measurement harness looks like this:

import torch

def measure_step(model, optimizer, batch):
    torch.cuda.reset_peak_memory_stats()

    optimizer.zero_grad(set_to_none=True)
    outputs = model(**batch)
    outputs.loss.backward()
    optimizer.step()

    torch.cuda.synchronize()

    allocated_gib = torch.cuda.max_memory_allocated() / (2 ** 30)
    reserved_gib = torch.cuda.max_memory_reserved() / (2 ** 30)
    return allocated_gib, reserved_gib

max_memory_allocated tracks active tensor allocations. max_memory_reserved includes blocks held by the CUDA allocator for reuse. A large reserved number with a much smaller allocated number points more toward caching or fragmentation than live activations. Calling empty_cache may change the display without fixing the underlying workload.

Next, vary one dimension at a time. Run micro-batch 1 and 2 at the same sequence length. Then hold the micro-batch constant and vary sequence length. If peak memory barely changes while model state is already near capacity, checkpointing will not rescue the run. If the peak rises substantially with tokens, activations are a credible target.

Then checkpoint the repeated Transformer blocks, not an arbitrary wrapper around the entire training step. Compare:

• peak allocated and reserved GiB;
• tokens per second;
• time spent in forward and backward;
• loss and gradient behavior on a fixed small batch;
• validation loss over a short run;
• communication time if the model is distributed.

If micro-batch 1 is the only configuration that fits, gradient accumulation may preserve the effective batch without increasing activation memory. With two accumulation steps, the effective batch is approximately B_micro × 2, while each individual forward and backward still sees B_micro. It costs extra passes and may not improve throughput, but it is a useful baseline against checkpointing.

For full fine-tuning, compare checkpointing with the model-state remedies as well. Mixed precision changes tensor storage and kernel behavior. Distributed training can shard or distribute parts of the state. Neither replaces activation checkpointing, but either may be the better first move.

One particularly useful distinction is when optimizer state appears. Standard PyTorch Adam normally allocates its moment state lazily at the first optimizer.step(), after the first backward pass. The first symptom usually tells you where to look:

• OOM before a meaningful forward pass: suspect parameters, inputs, initialization, or preallocated sharding state.
• OOM at the first optimizer.step(): suspect lazy optimizer-state allocation.
• OOM near the end of forward or at the start of backward, with the peak scaling with depth or tokens: suspect retained activations.
• OOM only at long sequences despite checkpointing: inspect attention workspaces and uncheckpointed kernels.
• Loss changes sharply after enabling it: inspect dropout RNG handling, mutable state, and checkpoint boundaries.
• Throughput collapses without allowing a larger micro-batch: remove it or use a more selective policy.

The deeper lesson is larger than this one PyTorch feature. A tensor does not have to remain resident just because backward will eventually need its information. You can store it, recompute it, compress it, shard it, or move it elsewhere. Each choice spends a different resource.

The activation checkpointing lesson goes through the replay schedule and configuration details. The right answer for a particular job still comes from the ledger and the profiler, not from a flag’s reputation.

Sources

• PyTorch checkpoint documentation
• PyTorch activation-checkpointing techniques
• Hugging Face gradient checkpointing