Gradient Checkpointing

Backpropagation through $L$ layers ordinarily keeps every intermediate activation alive until its gradient is consumed, giving $O(L)$ memory. Checkpointing keeps only every $s$-th activation and re-runs the forward pass within a segment when the backward pass reaches it.

The cost is one extra forward pass over the checkpointed region: about a third more compute for a transformer, since a forward pass is roughly half the cost of forward-plus-backward.

import torch
from torch.utils.checkpoint import checkpoint
from torch import Tensor, nn


class CheckpointedStack(nn.Module):
    def __init__(self, layers: list[nn.Module], every: int = 4):
        super().__init__()
        self.layers = nn.ModuleList(layers)
        self.every = every

    def forward(self, x: Tensor) -> Tensor:
        for i, layer in enumerate(self.layers):
            if self.training and i % self.every == 0:
                # use_reentrant=False replays RNG state correctly
                x = checkpoint(layer, x, use_reentrant=False)
            else:
                x = layer(x)
        return x

Leave use_reentrant at False. The reentrant implementation does not compose with anything that inspects the graph, silently drops gradients for inputs that do not require grad, and mishandles RNG in exactly the way described above.