Selective Recomputation

Gradient checkpointing is all-or-nothing per segment: keep the boundary, discard the interior, recompute it on the way back. With a whole transformer block as the segment, that is a second forward pass — about 36% more compute for the run.

The observation Korthikanti et al. make is that the memory is not evenly distributed over that block, and neither is the cost of recomputing it.

Recompute the low-intensity tensors: large, and nearly free to regenerate. Keep the high-intensity ones, where storing is cheaper than repeating the arithmetic.

Where each tensor falls

The quadratic tensors sit in the attention block and none of them costs much to regenerate — the scores need one matmul, and the softmax and dropout are elementwise. The linear tensors are matmul outputs, where recomputing means repeating the expensive part.

So: recompute the attention block, store everything else. Roughly 70% of full checkpointing’s saving at roughly a ninth of its overhead.

Composing it with the rest

Memory techniques stack, and the order they should be applied in is roughly the order of what they cost:

1. FlashAttention — free, removes the quadratic term.
2. Selective recomputation — ~4% compute.
3. Full recomputation on some layers — ~36% compute on those layers.
4. Offloading — bounded by PCIe.

Point 3 is worth stating separately because it is often missed: recomputation does not have to be uniform. Megatron supports checkpointing the first $k$ layers fully and the rest selectively, which gives a continuous dial between the two rather than a binary choice.

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


class Block(nn.Module):
    """Recompute the attention interior; keep the projections."""

    def __init__(self, dim: int, heads: int, selective: bool = True):
        super().__init__()
        self.n1, self.n2 = nn.RMSNorm(dim), nn.RMSNorm(dim)
        self.qkv = nn.Linear(dim, 3 * dim, bias=False)
        self.proj = nn.Linear(dim, dim, bias=False)
        self.ffn = FeedForward(dim)
        self.selective = selective

    def _core(self, qkv: Tensor) -> Tensor:
        """Scores, softmax, dropout, AV — large and cheap to redo."""
        return attention_core(qkv)

    def forward(self, x: Tensor) -> Tensor:
        qkv = self.qkv(self.n1(x))                  # kept: a large matmul

        core = checkpoint(self._core, qkv, use_reentrant=False) if self.selective \
            else self._core(qkv)

        x = x + self.proj(core)
        return x + self.ffn(self.n2(x))             # kept

use_reentrant=False is the modern implementation and the one to use. The reentrant version runs the recomputed forward outside autograd’s normal machinery, which breaks with keyword arguments, with tensors that do not require grad, and with anything stateful inside the segment — a long tail of failures that the non-reentrant version does not have.

The table above is for a standard transformer at typical shapes. It moves — with FlashAttention, with a very short sequence, with a wide FFN — and the honest approach is to measure rather than to trust it. torch.cuda.memory_allocated around each candidate segment, and the profiler’s memory timeline, will name the tensors that actually dominate on your configuration in a few minutes.