Gradient Checkpointing
Store activations at segment boundaries only, and recompute the interior during the backward pass. With √L segments both memory and recompute land at O(√L).
Standing
Current standardPart of the default recipe. Present in most frontier models trained today, and the thing a new design departs from rather than argues for.
judged as of 2026-09 · what the labels mean
Theory
Backpropagation through layers ordinarily keeps every intermediate activation alive until its gradient is consumed, giving memory. Checkpointing keeps only every -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.
Implementation
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 xLeave 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.