Activation checkpointing — why training can save memory by deliberately recomputing work

Recent lessons have focused heavily on inference: KV caches, quantization, PagedAttention, continuous batching, and the separation of prefill from decode.

Today we return to training, where the memory problem is different.

During inference, the model mostly needs its weights plus the state required to continue generation. During training, backpropagation needs information from the forward pass so that it can compute gradients in reverse.

That creates a surprisingly large memory bill.

The key idea for today is:

This is activation checkpointing, also called gradient checkpointing.

It deliberately performs extra computation in exchange for using less memory.

Why training remembers the forward pass¶

Consider a simple chain of layers:

The forward pass computes:

and so on.

Eventually we compute a loss:

Training then runs backward.

Using the chain rule:

But evaluating the derivative of \(f_4\) usually requires information from its forward computation.

The same is true as we continue backward through \(f_3,f_2,f_1\).

So the straightforward implementation keeps intermediate activations:

forward:

x0 → [f1] → x1 → [f2] → x2 → [f3] → x3 → [f4] → x4
             save          save          save

backward:

x0 ← [f1] ← x1 ← [f2] ← x2 ← [f3] ← x3 ← [f4] ← loss
             use           use           use

Those saved tensors are activations.

Why activations can become enormous¶

A model parameter is stored once per model replica.

Activations scale with the amount of data flowing through the model.

Very roughly, an activation tensor might have shape:

where:

• \(B\) is batch size,
• \(T\) is sequence length,
• \(d\) is hidden dimension.

If stored in 2-byte precision, one such tensor costs approximately:

bytes.

Suppose:

Then one tensor contains:

values.

At 2 bytes each:

That is for one tensor.

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

Not every implementation stores every full tensor, and modern kernels can fuse operations or avoid materializing some intermediates. But the scaling problem remains:

and training needs enough information to run backward.

The naive solution: save everything¶

Imagine a 12-layer network.

A simplified forward pass might save the output of every layer:

input
  ↓
L1  → save A1
  ↓
L2  → save A2
  ↓
L3  → save A3
  ↓
...
  ↓
L12 → save A12
  ↓
loss

When backward begins, the needed activations are already available.

This minimizes recomputation.

But it maximizes activation storage.

If each saved layer state costs \(M\), then the simplified memory cost scales like:

where \(L\) is the number of layers.

For a large Transformer, that can determine whether the training job fits at all.

The checkpointing idea¶

Instead of saving every layer's activations, save only selected checkpoints.

For example:

L1
 ↓
L2
 ↓
L3  → CHECKPOINT
 ↓
L4
 ↓
L5
 ↓
L6  → CHECKPOINT
 ↓
L7
 ↓
L8
 ↓
L9  → CHECKPOINT
 ↓
L10
 ↓
L11
 ↓
L12 → CHECKPOINT

During the original forward pass, activations inside each region can be discarded.

Later, when backward needs them, we recompute that region from its checkpoint.

So instead of paying:

to preserve every intermediate value, we pay additional:

to recreate some of them.

That is the trade.

A concrete four-layer example¶

Return to:

Suppose we checkpoint only:

During the first forward pass:

x0
 ↓
f1
 ↓
x1      discard after it is no longer needed
 ↓
f2
 ↓
x2      SAVE
 ↓
f3
 ↓
x3      discard
 ↓
f4
 ↓
x4      SAVE

Now backward starts at \(x_4\).

To backpropagate through \(f_4\) and \(f_3\), we need the missing forward intermediates around that region.

But we still have \(x_2\).

So we replay:

Conceptually:

saved x2
   ↓
re-run f3
   ↓
recreate x3
   ↓
re-run f4
   ↓
recreate what backward needs
   ↓
backprop through f4 and f3

Then those temporary recomputed activations can be freed.

To continue backward through \(f_2\) and \(f_1\), we similarly replay forward from \(x_0\).

We have traded storage for repeated forward computation.

The road-trip mental model¶

Imagine driving from Boston to San Francisco while writing an extremely detailed travel diary.

The naive strategy says:

At every mile, photograph the road, record your exact position, and keep every observation because you may need to reconstruct the trip backward later.

Your storage fills rapidly.

Checkpointing says:

Keep detailed records only at major cities.

For example:

Boston
   ↓
Chicago      CHECKPOINT
   ↓
Denver       CHECKPOINT
   ↓
Salt Lake    CHECKPOINT
   ↓
San Francisco

If you later need the details between Chicago and Denver, you retrace that segment.

You spend more driving time, but carry far less stored information.

The important question becomes:

How frequently should we place checkpoints?

Checkpoint frequency is a tradeoff¶

Checkpoint every layer:

L1* L2* L3* L4* L5* L6*

where * means saved.

Memory use is high, but recomputation is low.

Checkpoint rarely:

L1  L2  L3* L4  L5  L6*

Memory use falls, but more forward work must be replayed.

At the extreme, if we retained almost nothing, backward might require extensive recomputation.

So checkpoint placement balances:

This is one of the most important recurring patterns in large-model systems.

Why this is especially useful for Transformers¶

A Transformer consists of many repeated blocks:

Transformer block 1
Transformer block 2
Transformer block 3
...
Transformer block L

That gives us natural checkpoint boundaries.

A common conceptual strategy is to checkpoint the input to each Transformer block while allowing many internal intermediates to be discarded.

During backward, the block's forward computation is replayed to regenerate what autograd needs.

The exact strategy varies by framework and implementation. Some systems checkpoint whole blocks; others selectively preserve expensive-to-recompute values while recomputing cheaper operations.

The important point is that checkpointing need not be all-or-nothing.

Selective recomputation can be smarter¶

Consider two intermediate tensors.

Tensor A costs:

to store but is cheap to recompute.

Tensor B costs:

to store but requires an expensive operation to regenerate.

If memory is scarce, the sensible choice may be:

discard A → recompute later
save B    → avoid expensive replay

So the best policy is not simply:

Save every \(N\)th tensor.

It can depend on the ratio:

This is analogous to the cache-value reasoning from our prefix-caching lesson.

Different intermediates have different economic value.

Why FlashAttention fits the same philosophy¶

Earlier we studied FlashAttention.

Its important insight was that attention need not materialize the full:

attention matrix in high-bandwidth memory.

Instead, the computation is tiled so that intermediate values can be produced, consumed, and discarded efficiently.

Activation checkpointing has a related philosophy:

Sometimes recomputation is cheaper than memory traffic or storage.

Modern AI systems increasingly treat computation as relatively abundant and high-bandwidth memory as precious.

Training memory has several major residents¶

It helps to separate the memory bill into categories.

A simplified training process may need memory for:

model parameters
gradients
optimizer state
activations
temporary kernel workspace

Checkpointing primarily attacks:

It does not magically shrink the parameters.

It does not remove Adam's optimizer statistics.

It does not eliminate gradients.

This matters because the dominant memory term depends on the training setup.

Other techniques attack other terms.

For example:

mixed precision
→ smaller numerical representations

optimizer-state sharding
→ distribute optimizer memory

parameter sharding
→ distribute model weights

activation checkpointing
→ store fewer forward intermediates

Large-scale training usually combines several of these.

Why longer context makes checkpointing increasingly valuable¶

Recall that many activations scale with:

Double the sequence length \(T\), and many activation tensors roughly double in size.

Some attention-related intermediates can have even worse naive scaling with sequence length, which is why memory-efficient attention kernels matter so much.

So moving from:

to:

can dramatically increase training-memory pressure even though the model has exactly the same number of parameters.

This is an important distinction:

Batch size and context length matter enormously.

Checkpointing can enable a larger batch¶

Suppose a GPU has enough memory for:

without checkpointing.

After reducing activation storage, perhaps it can fit:

Now the training step performs more recomputation per example, but processes twice as many examples per step.

Whether this is faster overall depends on hardware utilization and communication.

This creates a subtle possibility:

A technique that adds FLOPs can sometimes improve end-to-end throughput if the saved memory enables a more efficient batch shape.

Extra arithmetic is not automatically slower if the original workload was constrained by memory or poor utilization.

Checkpointing interacts with distributed training¶

Suppose a model is spread across many GPUs.

Each GPU may own only part of the model, but it still needs activations for its local backward computation.

Activation checkpointing reduces how much of that state must remain resident between forward and backward.

This can make it easier to fit:

larger microbatches
longer sequences
larger models

But recomputation also means replaying operations on every participating device.

If those operations involve communication across tensor-parallel or sequence-parallel groups, the cost can include not only extra FLOPs but extra communication.

So the true cost is closer to:

Again, systems details matter.

Determinism creates a subtle requirement¶

Suppose a checkpointed region contains dropout.

The original forward pass samples a random mask:

If recomputation later uses a different mask:

then we have not actually recreated the same forward computation.

The gradients can be wrong.

Checkpointing systems therefore need to preserve or reconstruct relevant random-number-generator state so that recomputed stochastic operations behave consistently.

Conceptually:

original forward:
x → dropout(mask = m) → y

recomputed forward:
x → dropout(mask = m) → y

not:

x → dropout(mask = m') → different y

This is a useful reminder that "recompute" means reproduce the required computation, not merely run approximately similar code again.

Why inference usually does not need this trick¶

During ordinary inference, we do not run backward.

So there is no future gradient computation waiting to consume all the forward activations.

Most temporary activations can be discarded as soon as the next operation no longer needs them.

That is why activation checkpointing is primarily a training technique.

Inference has its own persistent-state problem:

This gives us a clean contrast:

training:
save forward state for backward
→ activation-memory pressure

autoregressive inference:
save past K/V for future tokens
→ KV-cache pressure

Both systems retain information from the past, but for completely different reasons.

The deeper idea: memory and compute are exchangeable resources¶

At first, a computation graph seems to imply a fixed recipe:

compute value
store value
use value later

Checkpointing breaks that assumption.

If a value can be deterministically regenerated from something smaller that we retained, then storage becomes optional.

We can choose:

or:

This is a form of time-space tradeoff.

Computer science is full of them.

A lookup table spends memory to save computation.

Compression spends computation to save storage.

Caching spends memory to avoid recomputation.

Activation checkpointing does the reverse:

Notice the symmetry with prefix caching from Day 50.

Prefix caching says:

Computation is expensive enough that we should spend memory to avoid repeating it.

Activation checkpointing says:

Memory is expensive enough that we should repeat computation instead of storing its result.

The correct direction depends on which resource is scarce.

The key mental model¶

Imagine every intermediate activation comes with two price tags:

KEEP ME:
cost = bytes of scarce GPU memory

THROW ME AWAY:
cost = FLOPs needed to recreate me later

Training does not have to choose the same answer for every tensor.

For some activations:

so discarding them is attractive.

For others:

so preserving them makes sense.

That gives us today's main principle:

This is one reason enormous models can be trained at all.

Large-scale training is not simply "run backpropagation on bigger GPUs."

It is an exercise in deciding, tensor by tensor and phase by phase, which scarce resource to spend: memory capacity, memory bandwidth, compute, or communication.

Activation checkpointing is one of the clearest examples of that systems-level bargain.