Activation Checkpointing

Activation checkpointing (also called gradient checkpointing) saves memory by not storing intermediate activations during forward pass — instead recomputing them during backward. It’s the most effective single technique for reducing activation memory.

1. The Core Tradeoff

	Without checkpointing	With checkpointing
Forward	Compute and save all activations	Compute, save only inputs
Backward	Use saved activations	Recompute activations from saved inputs, then use
Memory	$O(L)$ layers × all tensors	$O(L)$ layers × input tensor only
Compute	1× forward + 1× backward	2× forward + 1× backward (+33%)

For a 70B model: 224 GB → 10.7 GB activations (21× reduction) at the cost of 33% more compute.

2. Implementation

PyTorch

from torch.utils.checkpoint import checkpoint

class CheckpointedTransformerLayer(nn.Module):
    def __init__(self, layer):
        super().__init__()
        self.layer = layer
    
    def forward(self, x):
        return checkpoint(
            self.layer,
            x,
            use_reentrant=False,  # Always use non-reentrant (safer, composable)
        )

# Apply to all layers
for i, layer in enumerate(model.layers):
    model.layers[i] = CheckpointedTransformerLayer(layer)

With FSDP2

from torch.utils.checkpoint import checkpoint

for layer in model.layers:
    # Activation checkpointing first
    layer.forward = lambda x, l=layer: checkpoint(l._original_forward, x, use_reentrant=False)
    # Then FSDP
    fully_shard(layer)

With Megatron-LM

# In Megatron config
--recompute-activations           # full checkpointing
--recompute-granularity selective  # selective checkpointing
--recompute-method uniform        # checkpoint every N layers
--recompute-num-layers 1          # checkpoint every layer

3. Selective Checkpointing

Not all activations cost the same to save or recompute:

Tensor	Memory cost	Recompute cost	Worth checkpointing?
Attention weights ( $s \times s$ )	Very high ( $O(s^2)$ )	High (full attention recompute)	✅ Yes — biggest savings
FFN intermediate	High ( $b \times s \times d_{ff}$ )	Low (one GEMM)	✅ Yes
LayerNorm inputs	Low ( $b \times s \times h$ )	Very low	⚠️ Optional
QKV projections	Moderate	Low (one GEMM)	⚠️ Optional
Dropout masks	Low (1 bit per element)	Free (regenerate from seed)	❌ Skip — just save it

Selective strategy: Save cheap tensors (dropout masks, LayerNorm inputs), recompute expensive ones (attention weights, FFN intermediates).

Strategy	Memory reduction	Compute overhead	When to use
None	0%	0%	Small models that fit
Selective	~60-70%	~5-15%	Best throughput/memory tradeoff
Full (per-layer)	~95%	~33%	Large models, long sequences
Full (per-block, every 2 layers)	~90%	~25%	Compromise

FlashAttention already eliminates the $O(s^2)$ attention activation by recomputing attention during backward. With FlashAttention, selective checkpointing focuses on FFN intermediates.

4. Interaction with Other Techniques

Technique	With checkpointing	Effect
FlashAttention	Both reduce activations	Complementary — FA handles attention, checkpointing handles FFN
FSDP	FSDP reduces static, checkpointing reduces dynamic	Must use together for large models
Tensor Parallelism	TP reduces both	Checkpointing still helps for remaining per-GPU activations
Gradient accumulation	Accumulation doesn’t reduce peak activation	Checkpointing does — they’re orthogonal

5. When Checkpointing Is Mandatory

Model	Without checkpointing	With full checkpointing	Verdict
7B, TP=1, b=4, s=4K	~76 GB activations	~5 GB	Mandatory
7B, TP=8, b=4, s=4K	~10 GB activations	~0.7 GB	Optional (helps with batch size)
70B, TP=8, b=1, s=4K	~12 GB activations (FA)	~0.7 GB	Strongly recommended
70B, TP=8, b=1, s=128K	very large	~22 GB	Absolutely mandatory

For detailed activation memory formulas, see Activation Memory Analysis.

Self-Check

Explain the compute-memory tradeoff of activation checkpointing (33% compute for 21× memory)
Implement per-layer checkpointing in PyTorch with use_reentrant=False
Compare full vs selective checkpointing strategies
Explain why FlashAttention and checkpointing are complementary
Determine when checkpointing is mandatory vs optional for a given configuration