41. Gradient Checkpointing Across a Layer Stack

Problem

Training a deep stack of layers normally means storing every layer's input activation for use during backward — 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 also maximizes recompute overhead; in practice you checkpoint every N layers ("segments") to tune the memory/compute tradeoff to the model and hardware at hand.

Implement forward_with_checkpointing: run x through every layer in layers, but group the layers into consecutive segments of size segment_size and checkpoint each segment's forward pass via torch.utils.checkpoint.checkpoint, so activations inside a segment are recomputed during backward instead of kept in memory. The output — and every gradient computed from it — must be numerically identical to a plain uncheckpointed forward pass through all layers.

Source: src/41_gradient_checkpoint_stack.py

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

>>> 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. Write a small helper (or inline closure) that runs a *list* of layers sequentially on an input — this is the function that will get checkpointed, one call per segment.
2. Iterate over layers in strides of segment_size, slicing out each consecutive segment.
3. For each segment, call torch.utils.checkpoint.checkpoint(segment_fn, x, use_reentrant=False) instead of just running the segment directly.
4. Feed the output of each checkpointed segment as the input to the next segment, exactly as in a plain sequential forward pass.
5. Handle the last, possibly-shorter segment automatically: Python slicing (layers[start:start+segment_size]) naturally truncates at the list's end, so no special-casing is needed.
6. Return the final x after all segments have run.

The key insight is that checkpoint() is transparent to correctness by design: it re-runs the exact same forward function during backward to regenerate the activations it didn't keep, so the values it produces are bit-identical to not checkpointing at all — the only thing that changes is when the compute for those activations happens, and whether they're held in memory in the meantime.

Reference solution

import torch
import torch.nn as nn
import torch.utils.checkpoint as checkpoint


def forward_with_checkpointing(
    layers: nn.ModuleList, x: torch.Tensor, segment_size: int,
) -> torch.Tensor:
    def run_segment(segment_layers, inp):
        for layer in segment_layers:
            inp = layer(inp)
        return inp

    for start in range(0, len(layers), segment_size):
        segment = list(layers[start:start + segment_size])
        # use_reentrant=False is the modern, recommended checkpoint mode:
        # activations *inside* this segment are freed after the forward
        # and recomputed on demand during backward
        x = checkpoint.checkpoint(lambda inp, seg=segment: run_segment(seg, inp), x, use_reentrant=False)
    return x

Key Functions & Tricks

• torch.utils.checkpoint.checkpoint(fn, *args, use_reentrant=False) — wraps a function so its intermediate activations aren't stored; use_reentrant=False is the modern mode (supports non-tensor outputs, keyword args, and nested checkpointing more reliably than the legacy reentrant mode).
• List slicing with a default-argument closure (lambda inp, seg=segment: ...) — captures each loop iteration's segment by value, avoiding the classic Python late-binding-closure bug where every lambda would otherwise see the same, final segment.
• range(0, len(layers), segment_size) — strided iteration that naturally produces a shorter final segment when len(layers) isn't a multiple of segment_size, no extra edge-case code required.
• Comparing checkpointed vs. plain forward/backward on cloned, independently-tracked inputs (x.clone().detach().requires_grad_(True)) — the standard way to verify a memory-saving rewrite is numerically a no-op.

How to Recognize This Pattern