Vol. XIV · Parallelism Plate 2 of 11activation checkpointing, recomputation, rematerialisation, selective recomputation, memory–compute trade-offCE

Recomputing Activations.

Throw the activations away and recompute them during the backward pass — memory for compute, at a price you can choose.

Activations are stored only because the backward pass needs them. If memory is short, they can be discarded after the forward pass and recomputed from a saved checkpoint — the input to each layer — when the backward pass reaches that layer. With every layer checkpointed, the stored activations shrink to one vector per token per layer, at the cost of a second forward pass: training goes from about 6N to 8N FLOPs per token, a third more.

Recomputation can be selective. Korthikanti and colleagues noticed that much of a layer’s activation memory sits in the attention scores and softmax, which are cheap to recompute but large to store, while the matrix multiplies are expensive to recompute and small to store. Recomputing only the former saves most of the memory for a few per cent of extra compute — and FlashAttention (plate XVI.4) makes the same trade inside its kernel. The workshop compares the strategies for a Llama 3 configuration.

Below: left, activation memory per device for no recomputation, selective recomputation of attention, and full recomputation, for the selected model and tokens. Right, the extra compute each costs.

What you will learn
Workshop
Activation memory per device for the selected Llama 3 configuration, tokens per micro-batch as chosen, sequences of 8 192: none recomputed (34 d + 5 H T bytes per token per layer); selective (attention scores recomputed: 34 d); full (only each layer’s input kept: 2 d). Right: extra compute relative to 6N per token — about 33% for full recomputation; the selective figure is the attention-score FLOPs re-done.

Memory for compute, selectively

For Llama 3 70B with 8 192 tokens per micro-batch, keeping every activation takes 1.9 TB, dominated by the attention scores. Recomputing just the scores and softmax brings it to 182.5 GB, saving 90% for 2.4% more compute; recomputing everything brings it to 10.7 GB, saving 99% for 33% more. The selective strategy is almost free because the parts it discards are the cheapest to recompute.

Full recomputation still has its place when memory is the binding constraint — it is how very large models fit at all on some layouts — and it is why hardware FLOPs utilisation can exceed model FLOPs utilisation (plate XIII.6): PaLM’s 57.8% HFU against 46.2% MFU is mostly recomputation. In practice frameworks expose a dial between the two, and FlashAttention has made the selective case the default, since it never stores the scores in the first place.

Chen and colleagues showed in 2016 that checkpointing every √n layers of an n-layer network needs memory proportional to √n for one extra forward pass — the origin of “sublinear memory” training.

Terms

Activation checkpointing
Storing only some activations and recomputing the rest in the backward pass.
Selective recomputation
Recomputing only the activations that are cheap to compute and large to store.
Full recomputation
Recomputing every layer’s activations from its saved input.
Memory–compute trade-off
Spending extra FLOPs to need less memory.
Recomputation trades compute for activation memory; recomputing only the cheap, bulky attention parts saves most of the memory almost for free.

Sources

  1. Chen, Xu, Zhang and Guestrin, “Training Deep Nets with Sublinear Memory Cost”, arXiv:1604.06174, 2016.
  2. Korthikanti et al., “Reducing Activation Recomputation in Large Transformer Models”, MLSys 2023.