Activation checkpointing — trading extra compute for much less training memory

Yesterday we saw gradient accumulation separate the physical microbatch from the effective global batch.

That solves one memory problem: we do not need millions of training tokens resident on the GPU simultaneously.

But there is a harder case.

Suppose even a single microbatch is too large because the Transformer must remember all of the intermediate activations needed for backpropagation.

Then shrinking the batch further may not be enough.

The standard solution is activation checkpointing, also called gradient checkpointing.

Its core idea is beautifully simple:

Instead of saving every activation from the forward pass, training saves only selected checkpoints. During backpropagation, missing activations are recomputed from those checkpoints.

We spend more arithmetic to use less memory.

Why training needs activations at all¶

During inference, a layer can often compute an output and then discard many temporary values.

Training is different because backward needs information from forward.

Consider the tiny function:

Its derivative is:

If backward wants to compute:

it needs \(x\):

So forward cannot always simply forget \(x\).

A neural network contains enormous numbers of such dependencies.

A Transformer layer leaves breadcrumbs¶

Imagine a simplified Transformer block:

input
  ↓
RMSNorm
  ↓
Q, K, V projections
  ↓
attention scores
  ↓
softmax
  ↓
attention output
  ↓
residual add
  ↓
RMSNorm
  ↓
MLP
  ↓
residual add

Backward must compute gradients through all of these operations.

To do that efficiently, the framework normally stores useful intermediate tensors during forward.

Those stored tensors are called activations.

For a deep Transformer, the forward pass therefore leaves a long trail of breadcrumbs:

layer 1 activations  ─┐
layer 2 activations   │
layer 3 activations   │
...                   ├─ kept until backward
layer 79 activations  │
layer 80 activations ─┘

That trail can consume enormous amounts of HBM.

Why activations become huge¶

A basic hidden-state tensor has shape roughly:

Suppose:

The hidden state alone contains:

values.

At BF16:

That is roughly half a gigabyte for one hidden-state-sized tensor.

A Transformer layer creates several intermediate tensors, and a large model may contain dozens or hundreds of layers.

Not every tensor has that exact size, and optimized kernels avoid materializing some intermediates, but the scaling problem is clear:

and, without recomputation strategies, roughly with depth as well.

This is especially painful for long-context training.

The naive strategy: save everything useful¶

Suppose our network has six conceptual blocks:

x
↓
A
↓
B
↓
C
↓
D
↓
E
↓
F
↓
loss

During forward, ordinary training might retain:

x, A, B, C, D, E, F

so backward can immediately use them:

loss
↑
F
↑
E
↑
D
↑
C
↑
B
↑
A
↑
x

This minimizes recomputation.

But it maximizes activation storage.

The checkpointing strategy¶

Instead, suppose we save only:

x
C
F

and discard:

A
B
D
E

After forward, memory contains:

x -------- C -------- F

Now backward reaches the final segment.

To compute gradients through:

D → E → F

we no longer have D and E.

But we still have C.

So we rerun that portion of forward:

C
↓
recompute D
↓
recompute E
↓
recompute F

and then immediately backpropagate through it.

Later, to backpropagate through:

A → B → C

we restart from saved \(x\):

x
↓
recompute A
↓
recompute B
↓
recompute C

Then backward continues.

We traded:

for:

A concrete memory example¶

Imagine each block produces 1 GB of activations that backward would otherwise need.

With 12 blocks, saving everything requires roughly:

Now checkpoint only every fourth block:

block 0  ✓
block 1
block 2
block 3
block 4  ✓
block 5
block 6
block 7
block 8  ✓
block 9
block 10
block 11
block 12 ✓

Instead of retaining every block's full activation state for the entire forward pass, we retain a much smaller set of boundary states and reconstruct intermediate states when needed.

The exact memory reduction depends on what each block saves internally, but conceptually we have changed the requirement from:

remember the whole journey

to:

remember a few landmarks
and retrace the road between them

That is the key mental model.

Why recomputation is often a good trade¶

At first this sounds wasteful.

GPUs are expensive. Why deliberately repeat computation?

Because training frequently hits a memory-capacity wall before it hits a compute wall.

Suppose without checkpointing we can fit:

tokens.

With checkpointing we might fit:

or use a larger microbatch.

If the alternative is:

the training configuration does not fit at all

then spending extra FLOPs is an excellent bargain.

This is a recurring systems principle:

More precisely, recomputation may be cheaper than keeping a large tensor resident in scarce HBM for a long time.

The surprising asymmetry between forward and backward¶

A training step already performs more computation than inference because it needs backward.

Very roughly, for many dense neural-network operations, backward costs on the same order as — and often around twice — the forward computation.

So a crude mental model might be:

forward          ≈ 1 unit
backward         ≈ 2 units
-------------------------
training         ≈ 3 units

Activation checkpointing adds some forward recomputation.

If we recompute much of the forward once during backward, we might move toward something like:

original forward       ≈ 1
recomputed forward     ≈ 1
backward               ≈ 2
---------------------------
checkpointed training  ≈ 4

These numbers are only a mental model, not a universal law.

Different operations and checkpointing strategies have different costs.

But they show why a substantial memory reduction may cost a much smaller proportional increase in total training compute.

If training goes from roughly 3 units to 4, that is about:

or a 33% increase in this simplified example — potentially worthwhile if memory usage falls dramatically.

Checkpoint granularity creates a tradeoff¶

We do not have to choose between:

save everything

and:

save almost nothing.

We can choose how frequently to checkpoint.

More checkpoints¶

✓ layer 1
✓ layer 2
✓ layer 3
✓ layer 4
...

means:

but:

Fewer checkpoints¶

✓ layer 1
  layer 2
  layer 3
  layer 4
✓ layer 5
...

means:

but:

So checkpoint placement is an optimization problem.

Why this matters especially for long context¶

Recall that many activation tensors scale with:

Doubling sequence length from:

to:

roughly doubles many activation tensors.

Attention can introduce even nastier intermediate scaling if implemented naively.

This is one reason techniques such as FlashAttention and activation checkpointing complement each other.

FlashAttention avoids materializing huge attention matrices in HBM.

Activation checkpointing avoids retaining many other forward intermediates across the entire forward/backward interval.

They attack different forms of memory pressure.

Checkpointing and gradient accumulation solve different axes¶

Yesterday's gradient accumulation says:

Process fewer training examples simultaneously.

If:

we can reduce:

while increasing accumulation steps \(A\).

Activation checkpointing says:

For the examples that are resident, remember fewer intermediate tensors.

So:

gradient accumulation
→ reduces simultaneous examples/tokens

activation checkpointing
→ reduces stored state per simultaneous example/token

They combine naturally.

A training configuration might use:

microbatch = 1 sequence/GPU
gradient accumulation = 64 steps
activation checkpointing = enabled

and still produce a huge global training batch.

There is another trick: selective checkpointing¶

Not every operation has the same ratio of:

Suppose operation A creates a huge activation but is cheap to recompute.

That is an excellent checkpointing target.

Suppose operation B creates a small activation but requires an enormous matrix multiplication to reconstruct.

Saving B may be better.

So advanced systems can selectively choose what to retain.

The ideal strategy is roughly:

This turns checkpointing into a resource-allocation problem rather than a simple on/off switch.

But recomputation must reproduce the forward pass¶

There is a subtle correctness issue.

Suppose the forward pass contains randomness, such as dropout:

If we discard \(y\) and later recompute it with a different random mask, backward would no longer correspond to the original forward computation.

Checkpointing systems therefore need to preserve or reconstruct the relevant random-number-generator state so recomputation is consistent.

Conceptually:

original forward:
x → dropout mask #417 → y

recomputed forward:
x → dropout mask #417 → y

not:

x → completely different mask → different y

Modern frameworks handle this machinery, but it illustrates an important point:

Recompute means reproduce the necessary forward computation, not merely run something vaguely equivalent.

Memory is really about tensor lifetime¶

Here is the deeper systems insight.

A tensor's memory cost is not only:

It is also:

Consider a 2 GB activation created near the beginning of forward.

Without checkpointing, it might remain resident through:

dozens of later layers
loss computation
start of backward
dozens of backward operations

before finally being freed.

That is a long lifetime.

Checkpointing says:

Don't keep this 2 GB object alive for the entire trip. Delete it now and recreate it near the moment backward actually needs it.

This idea — shortening tensor lifetimes — is fundamental to memory optimization.

A useful picture of peak memory¶

Imagine training memory over time:

memory
  ^
  |                         /\
  |                       /    \
  |                     /        \
  |                   /            \
  |                 /                \
  |_______________/                    \____
  +------------------------------------------> time
             forward             backward

During forward, saved activations accumulate.

Near the transition to backward, many of them are simultaneously alive.

That can create the peak.

Checkpointing changes the shape:

memory
  ^
  |              /\       /\       /\
  |             /  \     /  \     /  \
  |            /    \___/    \___/    \
  |___________/                         \___
  +------------------------------------------> time
        forward     recompute/backward

The system repeatedly recreates smaller groups of activations rather than retaining everything simultaneously.

Peak memory falls, even though total computation rises.

Connect this to yesterday's lesson¶

We can now distinguish three different large objects in training.

Optimizer state¶

AdamW stores parameter-sized history:

Reduce this with techniques such as optimizer sharding or lower-precision optimizer states.

Global batch¶

Millions of tokens may contribute to one update.

Reduce simultaneous batch residency with:

Activations¶

Each resident token generates intermediate state through many layers.

Reduce this with:

So a simplified training-memory decomposition is:

Different techniques target different terms.

This is why there is no single "make the model fit" trick.

The deeper mental model: memory as a cache of past computation¶

The most powerful way to think about activations is not as mysterious neural-network state.

They are a cache.

During forward, we compute intermediate values and save them because we predict backward will need them later.

That is exactly what a cache does:

to avoid:

Activation checkpointing deliberately chooses a lower cache hit rate.

Instead of:

backward needs activation
→ retrieve saved activation

we sometimes do:

backward needs activation
→ activation absent
→ recompute from checkpoint

So the tradeoff is:

This same tradeoff appears everywhere in computer systems.

Caches use memory to save compute or latency.

Compression uses compute to save storage or bandwidth.

Checkpointing uses compute to save HBM.

Put the last three lessons together¶

We can now see why training a giant LLM is much more than running backpropagation.

AdamW asks us to retain large parameter-sized optimizer states.

Gradient accumulation lets a huge statistical batch exist across time rather than simultaneously.

Activation checkpointing lets the forward pass itself exist partly across time: instead of preserving every intermediate result, we regenerate pieces when backward needs them.

The pattern is:

can't fit the optimizer?
→ shard it across machines

can't fit the global batch?
→ spread it across microbatch steps

can't fit the activations?
→ spread their existence across time by recomputing them

That leads to a broader principle of large-scale AI systems:

Space means multiple GPUs or machines.

Time means accumulation or recomputation.

A Transformer training run that appears on paper to be one enormous forward-and-backward computation is, in reality, carefully decomposed so that only a fraction of its mathematical state physically exists at any instant.

Activation checkpointing is one of the clearest examples of that engineering philosophy:

Forget something now, because recomputing it later is cheaper than remembering it the whole time.