How does gradient checkpointing reduce GPU memory, and what is the trade-off?

How to think about it

Gradient checkpointing, more precisely activation checkpointing, reduces peak GPU memory by saving only selected intermediate activations during the forward pass and recomputing the discarded activations during the backward pass. The trade-off is extra forward computation: you use less VRAM, but training usually takes longer.

Why it works

Suppose a model is a chain of layers:

x_i = f_i(x_(i-1))

An activation is an intermediate tensor such as the output of a layer. During the backward pass, the phase that calculates gradients from the loss, autograd needs values from the forward pass. For example, the gradient of a matrix multiplication needs the relevant input matrix. If the framework kept none of those values, it could not calculate the gradients efficiently.

Ordinary training therefore saves many activations. In a deep Transformer, that can mean tensors for every layer, attention projection, residual connection, normalization, and feed-forward operation. With a long sequence or a large micro-batch, these tensors can consume more memory than the model weights themselves.

Checkpointing changes the storage policy. You divide the forward computation into regions and save only the boundary input for each region. The intermediate tensors inside a region are discarded after its forward pass.

Later, when backward reaches that region, the framework restores its boundary input and runs the region’s forward computation again. This time it records the temporary activations needed for the backward calculation. Once their gradients have been computed, those temporary tensors can be released.

So checkpointing does not compress a tensor. It avoids keeping the tensor alive for the entire training step.

The important word is peak. The total amount of computation is not reduced, and the model still needs parameters, gradients, optimizer state, and temporary kernel workspaces. Checkpointing reduces the largest simultaneous set of live activation tensors.

Memory category	Does checkpointing reduce it?	Reason
Intermediate activations	Yes	Most are discarded and later recomputed
Parameters	No	The model still needs its weights
Parameter gradients	No	Backward still produces the same gradients
Optimizer state	No	Adam moments or similar state are unchanged
Temporary workspaces	Usually not directly	These depend on the kernels and implementation

Common misconception: gradient checkpointing is not the same as saving a model checkpoint to disk. A model checkpoint is a persistent snapshot for recovery or later loading. Activation checkpointing is a per-training-step memory technique.

A concrete Transformer example

Consider a 24-block Transformer with:

• micro-batch size: 2
• sequence length: 2,048 tokens
• hidden size: 4,096
• activation dtype: bfloat16, using 2 bytes per element

One residual-stream tensor has:

2 × 2,048 × 4,096 × 2 = 33,554,432 bytes

That is exactly 32 MiB, ignoring framework metadata.

Now make a deliberately simplified estimate. Suppose each block leaves eight tensors of roughly that size live for backward. Saving them across 24 blocks costs:

24 × 8 × 32 MiB = 6,144 MiB

That is 6 GiB of activation storage before counting attention-specific tensors, larger feed-forward tensors, embeddings, output layers, CUDA allocator overhead, and kernel workspaces. Real models vary, but the arithmetic shows why sequence length and micro-batch size matter.

Now group the 24 blocks into six checkpointed regions of four blocks each. The framework might retain only the input boundary for each region, plus the initial input. That is seven residual-sized tensors:

7 × 32 MiB = 224 MiB

The model does not run the whole step using only 224 MiB. During backward, it temporarily reconstructs one four-block region at a time, so the active workspace for that region is still present. But it no longer keeps the internals of all 24 blocks alive simultaneously.

This is the central result: persistent activation memory falls sharply, while recomputation happens locally during backward.

How I would use it

In PyTorch, a simple block-level pattern looks like this:

from torch.utils.checkpoint import checkpoint

class Model(torch.nn.Module):
    def forward(self, x):
        x = self.embed(x)

        for block in self.blocks:
            x = checkpoint(block, x, use_reentrant=False)

        return self.head(x)

The exact placement depends on the architecture. Transformer blocks are common boundaries because each block contains many intermediate tensors, while its input and output have a predictable shape.

Checkpointing every block is a reasonable starting point, not a law of nature. Checkpointing larger groups reduces the number of saved boundaries, but increases the temporary memory needed while recomputing a group. Checkpointing smaller groups lowers the recomputation workspace, but retains more boundaries. The best choice depends on the model’s memory profile and the batch and sequence lengths used in production.

I would benchmark the same representative workload with and without checkpointing. I would record peak allocated memory, peak reserved memory, step time, and tokens per second after a short warm-up. A setting that saves 8 GiB but makes training 45 percent slower may be sensible for a one-off experiment and a poor choice for a month-long pretraining run.

Many Transformer libraries expose a model-level gradient-checkpointing switch. That is convenient, but I would still verify which modules it wraps and which checkpoint implementation the current library version uses.

The trade-off and when the textbook answer is wrong

The cost is recomputation. A checkpointed region normally runs once during the forward pass and again when backward needs its discarded activations. This adds approximately another forward execution for each checkpointed region. The total training slowdown is not a fixed percentage because backward itself has a different cost, and attention, communication, data loading, and kernel efficiency may dominate the step.

Checkpointing is most useful when activations are the reason for the memory limit. It can allow a larger micro-batch, a longer context, or a model that otherwise would not fit. Sometimes that improves overall throughput because a larger micro-batch reduces gradient-accumulation overhead or makes better use of the GPU.

It is less attractive when the job is already compute-bound and has plenty of free VRAM. Recomputing forward passes then buys little. It also cannot rescue an out-of-memory error caused mainly by optimizer state or model parameters. For that problem, techniques such as mixed precision, parameter sharding, optimizer-state sharding, quantization, or a smaller model address the relevant memory category.

Sequence length adds another nuance. Many activation tensors scale linearly with the number of tokens, so doubling the sequence length roughly doubles those tensors. A naive attention implementation may also materialize an attention matrix whose size scales with the square of sequence length. Memory-efficient attention kernels avoid storing that full matrix, so after using such a kernel, checkpointing the feed-forward blocks may save more memory than checkpointing attention itself.

Checkpointing can also expose correctness problems. A checkpointed function must behave consistently when run a second time. Mutable module state, custom random-number generation, or side effects such as appending tensors to a Python list can make the recomputed result differ from the original. Standard PyTorch checkpointing normally preserves random-number state for ordinary dropout, but custom randomness and device movement deserve testing.

Failure modes I would check first

Symptom: peak VRAM barely changes and the job still fails with an out-of-memory error. The likely causes are that the memory-heavy modules were not checkpointed, the model is storing tensors outside autograd, or the dominant allocation is parameters, gradients, or optimizer state. I would inspect a memory snapshot rather than assume activations are the culprit.

Symptom: tokens per second drops much more than expected. The checkpoint regions may be too large, the workload may already be compute-bound, or the implementation may be recomputing more than intended. I would compare per-step time and kernel traces, then try different boundaries.

Symptom: the loss or gradients change only after enabling checkpointing. I would look for stateful code and non-reproducible randomness inside the checkpointed function. A checkpoint should recompute the same mathematical forward operation, not perform a second unrelated state update.

What they will ask next

Does checkpointing reduce total training memory?
It reduces the activation portion of peak memory, not the memory for parameters, gradients, or optimizer state. If Adam state is the problem, activation checkpointing is aimed at the wrong target.

Is checkpointing better than gradient accumulation?
They solve different problems. Gradient accumulation lowers the micro-batch size, so each forward and backward pass needs fewer activations, but it requires multiple passes to build the same effective batch. Checkpointing keeps the micro-batch size and recomputes activations. They are often combined.

Does checkpointing reduce training time?
Usually no. It adds forward computation during backward. It can still improve end-to-end throughput if the saved memory allows a larger micro-batch, longer useful context, or avoids splitting the model across slower devices. The answer must come from a benchmark on the target hardware.

For a deeper implementation treatment, see Activation Checkpointing.

Say this in the interview: “Activation checkpointing lowers peak VRAM by saving only region boundaries and recomputing discarded intermediates during backward; the cost is extra forward compute, while parameter, gradient, and optimizer memory remain unchanged.”