gradient checkpointing
Gradient checkpointing is a training technique that discards most intermediate activations during a neural network's forward pass and recomputes them as needed during the backward pass, trading extra computation for significantly reduced memory usage.
```flowchart
{
"rows": [
{ "type": "nodes", "items": [
{ "title": "Training normally stores every intermediate activation", "sub": "needed later for computing gradients, but consumes heavy memory", "tone": "neutral" }
]},
{ "type": "arrow" },
{ "type": "group", "title": "Gradient checkpointing discards most activations after use", "items": [
{ "title": "Only a subset of checkpoints kept in memory", "sub": "discarded activations recomputed later when actually needed", "tone": "blue" }
]},
{ "type": "arrow" },
{ "type": "nodes", "items": [
{ "title": "Memory usage drops substantially during training", "sub": "extra recomputation cost traded for that memory savings", "tone": "green" }
]}
]
}
```
**Gradient checkpointing exists because training large neural networks in the standard way requires storing every layer's intermediate activation values simultaneously, and for sufficiently large models that memory demand can exceed what available hardware can provide.** Since most of those stored activations are only needed briefly during the backward pass to compute gradients, gradient checkpointing instead keeps just a sparse set of checkpoint activations in memory and recomputes the rest on demand by re-running the relevant forward pass segment when the backward pass actually needs them, substantially cutting peak memory usage at the cost of some additional computation.
```svg
```
```svg
```
| Aspect | Standard training | Gradient checkpointing |
|---|---|---|
| Peak memory usage | Higher | Substantially lower |
| Computation per step | Baseline | Increased, due to recomputation |
| Maximum trainable model or batch size | More limited by memory | Larger, given the same memory |
| Common use | When memory isn't the binding constraint | Training large models under memory pressure |
**Gradient checkpointing typically increases total training compute by a modest, predictable amount, often cited as roughly one additional forward pass worth of computation, in exchange for its memory savings.** Because recomputing discarded activations means effectively re-running portions of the forward pass a second time during the backward pass, gradient checkpointing introduces a measurable but generally manageable increase in total compute, which is usually a worthwhile tradeoff when memory, rather than raw compute time, is the limiting factor.
**Choosing which activations to keep as checkpoints versus which to discard is itself a design decision, since checkpoint placement affects the balance between memory savings and how much recomputation the backward pass ends up needing.** Because keeping more checkpoints reduces the amount of recomputation needed but also reduces the memory savings, and keeping fewer checkpoints does the reverse, practical gradient checkpointing implementations often place checkpoints strategically, such as once per major network block, to find a reasonable balance between the two competing goals.
**Gradient checkpointing has become an important enabling technique for training very large models, since it lets a fixed amount of hardware memory support a larger model or batch size than would otherwise be possible.** Because the memory required to train a model without checkpointing scales directly with model size and batch size, gradient checkpointing has become a standard tool for pushing past what would otherwise be a hard memory ceiling, letting researchers and engineers train larger models on the same hardware than memory constraints alone would allow.
Read gradient checkpointing through a note-taking lens: rather than writing down every single step of a long calculation, you jot down just a few key checkpoints along the way and redo the missing steps from the nearest checkpoint whenever you actually need them, trading a bit of extra work for far less paper.