Activation Checkpointing › The Checkpointing Mechanism

Activation checkpointing redefines which activations must survive the full forward-to-backward journey. Understanding the mechanism requires tracing through both the forward and backward passes in detail.

The algorithm proceeds in three phases:

Segmented forward pass. The network is divided into segments, typically one segment per transformer layer or one per block of layers. At the boundary between segments, the activation tensor is saved to memory. Within each segment, intermediate activations are computed normally but immediately discarded after use, rather than being held for the backward pass.

Standard backward pass (outer). Backpropagation proceeds from the output layer toward the input. When it reaches a segment boundary, it has the activation at that boundary available (because it was checkpointed), and it triggers a recomputation pass for that segment.

Local recomputation. For each segment being backpropagated through, a mini-forward pass is re-run from that segment's saved input activation to its output. This regenerates all the intermediate activations within the segment. Backpropagation through the segment can then proceed normally using these regenerated activations, after which they are discarded again.

The memory at any point in backward is bounded by the number of checkpoint boundaries plus the memory for one segment's activations (the one currently being recomputed). If you have $L$ layers total and checkpoint every $k$ layers, you store $L/k$ checkpoint boundary tensors and hold at most $k$ layers of intermediate activations in memory at once during recomputation.

The behavior is like a rolling window: you always have the boundary tensors pinned, and the intermediate activations within one segment are temporarily materialized, then released as you move to the next segment. The maximum memory footprint is not $L$ times a layer's activation memory, but rather $L/k + k$ times a layer's activation memory.

Let's work through the memory savings concretely. With $L$ layers total, storing all activations requires memory proportional to $L$. Without checkpointing, this scales as $O(L)$ because every layer's activations coexist in memory simultaneously during the backward pass.

With checkpointing at every $k$ layers, the memory for activations is:

$$\text{Memory}_{\text{checkpointed}} = \frac{L}{k} \cdot M_{\text{ckpt}} + k \cdot M_{\text{layer}}$$

where:

• $L/k$ is the number of checkpoint boundaries stored simultaneously
• $M_{\text{ckpt}}$ is the memory for one checkpoint boundary tensor (the residual stream at a layer boundary)
• $k$ is the number of layers between consecutive checkpoints
• $M_{\text{layer}}$ is the activation memory for one full layer (the cost paid during recomputation)

The first term represents the checkpoint storage cost (boundaries stored during the entire forward pass), and the second term represents the recomputation cost (intermediate activations held during one backward segment pass).

To find the optimal $k$ that minimizes total memory, take the derivative with respect to $k$ and set it to zero:

$$\frac{d}{dk}\left[\frac{L \cdot M_{\text{ckpt}}}{k} + k \cdot M_{\text{layer}}\right] = -\frac{L \cdot M_{\text{ckpt}}}{k^2} + M_{\text{layer}} = 0$$

Solving for $k$:

$$k^* = \sqrt{\frac{L \cdot M_{\text{ckpt}}}{M_{\text{layer}}}}$$

When $M_{\text{ckpt}} \approx M_{\text{layer}}$ (checkpoint tensors are similar in size to per-layer activations), the optimal is $k^* = \sqrt{L}$, giving total memory proportional to $\sqrt{L}$ rather than $L$. For 32 layers and equal-sized tensors, the reduction is about 2.8x; the improvement can be larger when checkpoint boundaries are smaller than full-layer activations. The formula says something intuitive: you want to make the two terms in the memory expression roughly equal in magnitude, balancing the cost of storing boundaries against the cost of holding one segment in memory during recomputation.

However, in practice the $M_{\text{ckpt}} \approx M_{\text{layer}}$ assumption is often violated. A checkpoint boundary tensor is just the residual stream at one layer's output, with shape $[B, T, d]$. The full per-layer activation memory includes the attention weight matrix ($[B, H, T, T]$), which can be many times larger. This makes $M_{\text{ckpt}} \ll M_{\text{layer}}$, shifting the optimal $k$ to smaller values and making denser checkpointing more attractive.

The visualization below shows how activation memory scales with model depth under different checkpointing strategies. The baseline and per-layer checkpointing curves both grow linearly, but per-layer checkpointing has a much smaller slope because it stores only boundary tensors. Spacing checkpoints according to the analytical optimum grows as $O\left(\sqrt{L}\right)$ and produces the lowest activation-memory curve under the stated assumptions.

The recomputation adds a second forward pass per segment during backpropagation. In a standard training step, the relative compute cost is:

$$\text{Cost}_{\text{checkpointed}} = C_{\text{fwd}} + C_{\text{recompute}} + C_{\text{bwd}}$$

where:

• $C_{\text{fwd}}$ is the cost of the forward pass
• $C_{\text{recompute}}$ is the cost of the recomputed forward pass during backward (equal to $C_{\text{fwd}}$ for full per-layer checkpointing)
• $C_{\text{bwd}}$ is the cost of the backward pass, approximately $2 \times C_{\text{fwd}}$ (it must compute gradients for both weights and inputs at each layer)

For full per-layer checkpointing:

$$\begin{aligned} \text{Cost}_{\text{baseline}} &= C_{\text{fwd}} + 2 C_{\text{fwd}} = 3 C_{\text{fwd}} \\ \text{Cost}_{\text{checkpointed}} &= C_{\text{fwd}} + C_{\text{fwd}} + 2 C_{\text{fwd}} = 4 C_{\text{fwd}} \end{aligned}$$

The overhead ratio is $4/3 \approx 1.33$, confirming the often-cited "33% more compute" figure. This trade reduces activation memory by 10-20x for 33% more compute, which in practice means you can double or triple the batch size, fit a much larger model on the same hardware, or increase the sequence length substantially.

Why is the overhead only 33% and not 100%? Because the backward pass already costs twice the forward pass (to compute weight gradients and input gradients). The recomputation adds one forward pass on top of a baseline cost that already included two forward passes worth of work. If the backward pass were free, the overhead would be 100%. The expensive backward pass is what makes recomputation relatively cheap by comparison.

Selective checkpointing, covered below, reduces the 33% overhead by only recomputing the most expensive layers and leaving cheap activations in memory.

Let's trace through exactly what happens during a two-layer forward and backward pass with checkpointing at the layer boundary. This makes the memory states concrete.

Setup: Two transformer layers $L_1$ and $L_2$ with a checkpoint at their boundary. The input to $L_1$ is $x_0$, the output of $L_1$ (and input to $L_2$) is $x_1$, and the output of $L_2$ is $x_2$.

Forward pass, without checkpointing:

1. Compute $x_1 = L_1(x_0)$. Store all of $L_1$'s internal activations (attention weights, FFN intermediate, layer norm statistics).
2. Compute $x_2 = L_2(x_1)$. Store all of $L_2$'s internal activations.
3. Compute loss $\mathcal{L}(x_2)$.

At this point, both $L_1$'s and $L_2$'s internal activations are in memory simultaneously.

Backward pass, without checkpointing:

1. Compute $\partial \mathcal{L} / \partial x_2$ from the loss.
2. Use $L_2$'s stored activations to compute $\partial \mathcal{L} / \partial x_1$ and the weight gradients for $L_2$'s parameters. Discard $L_2$'s activations.
3. Use $L_1$'s stored activations to compute $\partial \mathcal{L} / \partial x_0$ and weight gradients for $L_1$'s parameters. Discard $L_1$'s activations.

Peak memory: $\text{Memory}(x_0) + \text{Memory}(L_1 \text{ internals}) + \text{Memory}(L_2 \text{ internals}) + \text{Memory}(x_1) + \text{Memory}(x_2)$.

Forward pass, with checkpointing at $x_1$:

1. Compute $x_1 = L_1(x_0)$. Save $x_1$ as a checkpoint. Discard $L_1$'s internal activations.
2. Compute $x_2 = L_2(x_1)$. Store $L_2$'s internal activations.
3. Compute loss $\mathcal{L}(x_2)$.

At this point, only $x_1$ (the checkpoint) and $L_2$'s activations are in memory. $L_1$'s internals have been discarded.

Backward pass, with checkpointing:

1. Compute $\partial \mathcal{L} / \partial x_2$ from the loss.
2. Use $L_2$'s stored activations to compute $\partial \mathcal{L} / \partial x_1$ and weight gradients for $L_2$'s parameters. Discard $L_2$'s activations.
3. Re-run $L_1$'s forward pass starting from the saved $x_0$ to reconstruct $L_1$'s internal activations. Discard the checkpoint $x_1$.
4. Use the recomputed $L_1$ activations to compute $\partial \mathcal{L} / \partial x_0$ and weight gradients for $L_1$. Discard $L_1$'s activations.

Peak memory: $\text{Memory}(x_0) + \text{Memory}(x_1) + \text{Memory}(L_2 \text{ internals})$ during step 5, and $\text{Memory}(x_0) + \text{Memory}(L_1 \text{ internals})$ during steps 6-7. $L_1$ and $L_2$ activations never coexist.

This is the fundamental gain: layers 1 and 2 no longer need to coexist in memory. For 32 layers, this prevents all 32 layers' activations from coexisting. Instead, you hold at most one layer's worth of internal activations at a time during recomputation.