Foundation Lab

Activation Checkpointing & Memory Efficiency

Essential for training large models—without it, you can't fit 100B models in GPU memory

Concept 79 of 100EfficiencyPhase 6
#79CheckpointingEfficiency
key equation
Memory: O(L) vs O(L)\text{Memory: } O(\sqrt{L}) \text{ vs } O(L)
Reading map and next steps

Selected Foundation Object

Keep the equation fixed; move through the evidence.

Concept 79 of 100CheckpointingEfficiency / Phase 6: Modern efficiency & inference
Current question

Backprop needs activations from forward pass; normally we store all of them

Memory: O(L) vs O(L)\text{Memory: } O(\sqrt{L}) \text{ vs } O(L)
PredictionCommit before tracing the equation.

Ask what should change under a concrete input, then trace that expectation through the equation.

EvidenceCompare the equation and source.

Use the key equation and canonical papers as the available witnesses, without implying that a runnable panel exists.

InvariantName what survives notation changes.

The useful learning product is the reusable mechanism you can carry into another model, paper, or engineering tradeoff.

Next moveContinue through the atlas.

Use prerequisites, dependents, and semantic links to repair the next gap without leaving the object behind.

Why It Matters for Modern Models

  • Essential for training large models—without it, you can't fit 100B models in GPU memory
  • Every major training framework (PyTorch, JAX) uses this technique
  • Memory-compute trade-off is fundamental: pay with one to save the other

What Tutorials Skip

What is still poorly explained in textbooks and papers:

  • Backprop needs activations from forward pass; normally we store all of them
  • Checkpointing says: "just save some, recompute the rest when needed"
  • Optimal checkpoint spacing is √L layers—minimizes memory × compute product

Visualization Status

Core Math (Optional Deep Dive)

If you want intuition first, start with the key equation and cited sources. Come back here for the full walkthrough.

Key Equation
Memory: O(L) vs O(L)\text{Memory: } O(\sqrt{L}) \text{ vs } O(L)

Memory problem: Storing activations for backprop requires O(L⋅B⋅d)O(L \cdot B \cdot d) memory.

Gradient checkpointing: Recompute instead of store:

  • Forward: Save only checkpoint activations (every L\sqrt{L} layers)
  • Backward: Recompute activations from nearest checkpoint

Memory: O(L)O(\sqrt{L}) instead of O(L)O(L)
Compute: 33% overhead (recompute forward pass once)

Selective checkpointing: Only checkpoint expensive layers (attention).

Trade-off: Memory×Compute≥constant\text{Memory} \times \text{Compute} \geq \text{constant}

Canonical Papers

Training Deep Nets with Sublinear Memory Cost

Chen et al.2016arXiv
Read paper →

Reducing Activation Recomputation in Large Transformer Models

Korthikanti et al.2022MLSys
Read paper →

Connections

Next Moves

Choose the next question to carry this object forward.