Activation checkpointing

Read it in Chapter 4 · checkpoint +50 XP →

Throw away activations during forward and recompute them in backward: trade ~33% more compute for far less memory.

Backpropagation needs the activations from the forward pass. For a transformer layer with FlashAttention that's roughly 34 · s · b · h bytes (sequence × micro-batch × hidden), and it adds up fast. A 7B model at micro-batch 4 needs over 70 GB of activations.

Checkpointing (also called recomputation) keeps only each layer's input, then re-runs the layer's forward during backward. Cost: an extra forward pass, ~33% more compute (8N instead of 6N FLOPs per token). Benefit: activation memory drops to about 2 · s · b · h bytes per layer.

Use it when activations are what's stopping you from fitting the model or a larger micro-batch. Turn it off when memory is fine. It's pure overhead then. Selective checkpointing (only the attention internals) gets most of the savings for a fraction of the cost.

Note that MFU deliberately doesn't count recomputed FLOPs, so checkpointing lowers MFU even though the GPU is just as busy.

Use it in a mission

Related

Sources & further reading

• Reducing Activation Recomputation (Korthikanti et al., 2022)
• Distributed Training (T. Tsugriy) — parallelism strategies, collectives, memory

Written for LazyLab in our own words; sources are cited for depth and verification.