Activation Checkpointing

The Memory Hog: Activations

When training a neural network, the forward pass computes outputs layer by layer. However, you can't just delete these outputs once the next layer is done.

backpropagation requires these intermediate values (activations) to calculate the gradients. Therefore, standard deep learning frameworks save every single activation from the forward pass in GPU memory until the backward pass reaches that layer.

For large models with large batch sizes, this Activation Memory dwarfs the memory taken up by the model weights. If you run out of memory during training, it is almost always because the activations filled up your VRAM.

The Compute/Memory Tradeoff

Activation Checkpointing (also called Gradient Checkpointing) solves this by recognizing a fundamental tradeoff: you can always trade compute for memory.

Instead of saving every activation, the framework only saves strategic "checkpoints" (e.g., every 4th layer).

• All the activations between the checkpoints are immediately deleted.
• This massively reduces the memory footprint during the forward pass.
• When the backward pass needs those deleted activations, it temporarily pauses, takes the nearest saved checkpoint, and re-runs the forward pass for that small block of layers to regenerate the missing values.

The Cost of Recomputation

Because Activation Checkpointing requires running parts of the forward pass twice, it makes training slower.

However, the overhead is surprisingly small. A full forward pass is generally cheaper than a backward pass. Checkpointing typically increases training time by about 20% to 30%.

In exchange, you reduce your activation memory footprint from $O(N)$ to roughly $O(\sqrt{N})$, where $N$ is the number of layers.

Why it Enables Massive Scale

While 30% slower sounds bad, Activation Checkpointing is often the only way to train large models.

By slashing memory requirements, checkpointing allows you to:

1. Fit the model onto the GPU in the first place.
2. Drastically increase your Batch Size.

Increasing the batch size improves GPU utilisation and mathematical efficiency. Often, the speed gained by processing a larger batch size completely eclipses the 30% compute penalty of recomputation, making your total training time faster.

Selective Checkpointing

Modern frameworks (like Megatron-LM and PyTorch) now use Selective Checkpointing.

Not all activations are created equal. Some operations (like the large matrix multiplications in the feed-forward networks) take a huge amount of compute but produce relatively small activations. Other operations (like the attention softmax) take very little compute but produce massive activation tensors.

Selective Checkpointing analyzes the computational graph and only checkpoints the operations that are computationally cheap but memory-heavy (like attention). This provides 90% of the memory savings with almost no compute penalty.

The Quick Version

• Standard training saves every intermediate activation in memory, causing Out of Memory (OOM) errors on large models.
• Activation Checkpointing deletes most activations and recomputes them on the fly during the backward pass.
• It trades a ~20% increase in compute time for massive memory savings.
• The memory savings allow you to use larger batch sizes, often resulting in faster overall training times despite the recomputation overhead.