Activation Checkpointing › The Memory Problem Activation Checkpointing Solves

To understand why activation checkpointing is necessary, you first need a precise picture of where memory goes during training. GPU memory is a scarce, non-expandable resource. Unlike CPU RAM, you cannot simply add more of it mid-run. When an allocation fails, the entire training process crashes, wasting potentially hours of compute time accumulated to that point.

As we examined in the Memory Management chapter, total GPU memory during training decomposes into four categories: model parameters, optimizer states, gradients, and activations. For large models with large batch sizes or long sequences, activations dominate.

The reason activations are so expensive comes down to the structure of the computation graph. Neural network training uses automatic differentiation, which traces through every operation performed during the forward pass and records it in a graph structure. Later, during backpropagation, the system walks this graph in reverse to compute gradients. For the gradient computation at each node to work, the system needs to know the values that existed at that node during the forward pass. So every intermediate tensor produced during the forward pass gets retained in memory until its gradient has been computed during the backward pass. For a deep network with many layers, "until the backward pass computes its gradient" effectively means "until the very end of the backward pass," because gradients flow from the output layer back to the first layer.

Consider a single transformer layer with hidden dimension $d$, sequence length $T$, and batch size $B$. The inputs and outputs of each sublayer (attention, feed-forward, layer norm) all need to be stored for the backward pass. The activations for a single transformer layer scale approximately as:

$$\text{Memory}_{\text{layer}} \approx 12 \cdot B \cdot T \cdot d \cdot \text{bytes}$$

where:

• $B$ is the batch size (number of sequences processed in parallel)
• $T$ is the sequence length (tokens per sequence)
• $d$ is the model hidden dimension
• the factor of 12 accounts for storing query, key, and value matrices, attention weight matrices, intermediate feed-forward activations, and residual stream tensors, each contributing a $B \times T \times d$ or $B \times T \times 4d$ tensor

This factor of 12 deserves more explanation because it surprises many practitioners who expect activations to be small. Within the attention sublayer alone, you store: the pre-norm input (one $[B, T, d]$ tensor), the query projection $Q$ ($[B, T, d]$), the key projection $K$ ($[B, T, d]$), the value projection $V$ ($[B, T, d]$), and the attention weight matrix $A = \text{softmax}(QK^T / \sqrt{d_k})$ ($[B, H, T, T]$). The attention weights alone contribute a tensor proportional to $T^2$, which is why long-context training is so memory-intensive. Within the feed-forward sublayer, the intermediate activation has shape $[B, T, 4d]$, four times the hidden dimension width.

For a 32-layer model with $d = 4096$, $T = 2048$, $B = 4$, and using float16 (2 bytes per value), total activation memory across all layers is:

$$\begin{aligned} \text{Memory}_{\text{total}} &= 32 \times 12 \times B \times T \times d \times 2 \text{ bytes} \\ &= 32 \times 12 \times 4 \times 2048 \times 4096 \times 2 \\ &\approx 64 \text{ GB} \end{aligned}$$

That is 64 GB just for activations, on a GPU that might have 40-80 GB total. The parameters themselves (at 16-bit precision) occupy roughly 14 GB for a 7B-parameter model, and the Adam optimizer states (maintained at 32-bit precision for numerical stability) add another 56 GB. You can see immediately why activation memory is not a footnote but the primary driver of training memory constraints. Without some intervention, a 7B model trained with Adam simply cannot fit in a single 80 GB A100.

Backpropagation requires activations because gradient computation is not invertible in general. To compute gradients with respect to the inputs of a layer, you must know what values those inputs had during the forward pass.

The chain rule makes this explicit. The gradient flowing back through a layer depends on the activation at that layer's input or output. For a nonlinear activation function $f$, the gradient $\partial \mathcal{L} / \partial x$ at some position $x$ requires knowing the value of $x$ itself:

$$\frac{\partial \mathcal{L}}{\partial x} = \frac{\partial \mathcal{L}}{\partial f(x)} \cdot f'(x)$$

where:

• $\mathcal{L}$ is the training loss
• $f'(x)$ is the derivative of the activation function, which depends on $x$ directly, not just on $f(x)$

For the GELU activation function specifically, $f'(x)$ has no simple closed form in terms of $f(x)$ alone, so $x$ must be retained separately from $f(x)$. You cannot reconstruct $x$ from $f(x)$ without inverting GELU, which is expensive and numerically unstable.

The same dependency appears in every major component of a transformer. The attention mechanism requires storing the attention weight matrix $A = \text{softmax}(QK^T / \sqrt{d_k})$ to compute gradient updates for the query, key, and value projection matrices $W_Q$, $W_K$, and $W_V$. The gradient of the softmax output with respect to its input is a Jacobian matrix, and computing it requires knowing the softmax output values themselves. The layer normalization sublayer must store its input mean and variance to correctly backpropagate through the normalization step.

Every nonlinear operation creates an activation dependency in the computation graph that must be resolved before the backward pass can continue. Without some form of checkpointing, the entire forward pass activation buffer must remain in memory until the corresponding backward pass gradient is computed, meaning activations from layer 1 must survive in GPU memory until after layer 32's backward pass completes.

It is helpful to think of training as maintaining a ledger of resources. On one side you have GPU memory capacity. On the other side you have compute budget in terms of FLOPs per second. The standard forward-backward cycle uses memory inefficiently: it reserves activation memory during the forward pass, holds it idle while the backward pass works layer by layer from the end, and only releases memory as each layer's gradients are computed.

This idle reservation is the waste that activation checkpointing targets. The forward pass activations from early layers sit in memory for the entire duration of the backward pass through later layers. If you have 32 layers, layer 1's activations are reserved while layers 32, 31, 30... 3, 2 all compute their gradients, finally being used only when the backward pass reaches layer 1. For a 32-layer model, early-layer activations are idle for roughly 97% of the backward pass. Discarding them and recomputing on demand turns idle memory into compute work, which is a much more efficient use of the available hardware budget.