45. Gradient Checkpointing Across a Layer Stack

Problem

Training deep stacks of layers means storing every layer's activations for the backward pass — memory that scales linearly with depth. Gradient (activation) checkpointing trades compute for memory: instead of storing every intermediate activation, you store only a handful of "checkpoints" and recompute the forward pass between them during backward. Checkpointing every single layer maximizes memory savings but maximizes recompute overhead too; in practice you checkpoint every N layers ("segments") to tune the tradeoff to the model and hardware at hand.

Implement a forward pass over a stack of layers that checkpoints each consecutive segment of segment_size layers, and must be numerically identical to a plain (non-checkpointed) forward pass.

Source: src/45_gradient_checkpoint_stack.py

def forward_with_checkpointing(
    layers: torch.nn.ModuleList,
    x: torch.Tensor,
    segment_size: int,
) -> torch.Tensor:
    ...

Examples:
>>> layers = nn.ModuleList([nn.Linear(4, 4) for _ in range(6)])
>>> x = torch.randn(2, 4, requires_grad=True)
>>> out = forward_with_checkpointing(layers, x, segment_size=2)
>>> out.shape
torch.Size([2, 4])

Step-by-Step Approach

1. Convert the ModuleList to a plain Python list so it can be sliced into segments.
2. Walk the layer list in steps of segment_size, slicing out one segment (a small sub-list of layers) per iteration.
3. For each segment, define a small closure that runs just those layers forward in order and returns the result — bind the segment by default-argument value, not by closing over the loop variable, so each closure captures its own segment correctly.
4. Pass that closure and the running x into torch.utils.checkpoint.checkpoint(...), and reassign x to its output before moving to the next segment.
5. Use use_reentrant=False — the modern non-reentrant checkpoint implementation, which handles arbitrary Python control flow inside the segment and doesn't require the input to already have requires_grad=True.
6. Return the final x after all segments have run.

The key insight is that checkpoint() doesn't change what gets computed — the forward output and every gradient are bit-for-bit equivalent to the non-checkpointed version. It only changes when the intermediate activations inside a segment exist in memory: discarded right after the forward pass, and rebuilt from scratch (a second forward call) only if backward actually needs them.

Reference solution

from torch.utils.checkpoint import checkpoint

def forward_with_checkpointing(layers, x, segment_size):
    layer_list = list(layers)
    for start in range(0, len(layer_list), segment_size):
        segment = layer_list[start : start + segment_size]

        def run_segment(inp, segment=segment):
            # bound by default arg, not the loop variable
            for layer in segment:
                inp = layer(inp)
            return inp

        # use_reentrant=False: supports arbitrary control flow, no
        # requires_grad restriction on inputs
        x = checkpoint(run_segment, x, use_reentrant=False)
    return x

Key Functions & Tricks

• torch.utils.checkpoint.checkpoint — wraps a function so its forward activations aren't stored; recomputes them from the segment's input during backward instead.
• use_reentrant=False — the recommended non-reentrant mode; avoids the older reentrant implementation's restrictions (no in-place ops issues, works with non-tensor inputs in the closure).
• Default-argument closure capture (segment=segment) — the standard Python fix for the late-binding closure bug in a loop.
• nn.ModuleList iteration/slicing — converting to a plain list first makes standard Python slicing (layer_list[start:start+segment_size]) available.

How to Recognize This Pattern