Activation Checkpointing & Memory Restructuring¶
Overview¶
Deep models die from activation memory, not parameters: every layer keeps its activations for backward. Activation checkpointing (a.k.a. gradient checkpointing) trades extra forward passes for memory: instead of storing every activation, store a subset and recompute the rest during backward.
- Memory for activations ∝ batch × depth × feature size.
- Checkpointing re-runs forward for the checkpointed segments during backward → ~O(sqrt(N)) memory from O(N).
- PyTorch:
torch.utils.checkpoint.checkpoint(fn, *args, use_reentrant=False). - Works hand-in-hand with FSDP (Ch 05-02), AMP,
torch.compile.
Golden rule: checkpoint where activations are big and cheap to recompute (feed-forward blocks, attention with large seq len); never checkpoint your critical small ops.
The Classic¶
import torch, torch.nn as nn
from torch.utils.checkpoint import checkpoint
class Block(nn.Module):
def __init__(self, d):
super().__init__()
self.fc1 = nn.Linear(d, d)
self.fc2 = nn.Linear(d, d)
def forward(self, x):
return self.fc2(torch.relu(self.fc1(x)))
model = nn.Sequential(*[Block(128) for _ in range(24)]).cuda()
# wrap each block in checkpointing in the forward:
class Checkpointed(nn.Module):
def __init__(self, blocks):
super().__init__()
self.blocks = blocks
def forward(self, x):
for b in self.blocks:
x = checkpoint(b, x, use_reentrant=False)
return x
model_ckpt = Checkpointed(model).cuda()
x = torch.randn(32, 128, device='cuda')
out = model_ckpt(x)
out.mean().backward()
print("OK:", out.shape)
How Much Memory Does It Save?¶
| Depth (blocks) | naive activations | checkpointed | peak ratio |
|---|---|---|---|
| 24 | 24 × block_acts | ~√24 × block_acts | ~4-5x |
| 100 | 100 × block_acts | ~10 × block_acts | ~10x |
The saving is on saved tensors kept for backward; parameters are unchanged.
use_reentrant: True vs False (important!)¶
| Mode | Behavior | When |
|---|---|---|
use_reentrant=False (new) |
no reentrant autograd trick; no double-backward hassle; supports nn.Module and arbitrary args | prefer this |
use_reentrant=True (legacy) |
reentrant checkpointing; needed for extreme memory; has restrictions (no views/in-place) | legacy/compat |
checkpoint(block, x, use_reentrant=False) # modern
checkpoint(block, x) # default tries reentrant; deprecated path
Granularity— Where to Cut¶
# Good
class TransformerBlock(nn.Module):
def forward(self, x, mask=None):
return checkpoint(self._inner, x, mask, use_reentrant=False)
def _inner(self, x, mask):
return self.mlp(self.attn(x, mask) + x) + x
Rules of thumb:
- Checkpoint per transformer block (not per single matmul) → ~2x compute, ~big memory drop.
- With BN: checkpoint segments must be big enough to amortize extra stat recompute.
- Don't checkpoint: embedding lookups, tiny layers, loss functions.
-
Pairs With (the full memory playbook)¶
- FSDP— FSDP shards params; checkpointing shards activations: the standard LLM combo.
- AMP— activations in fp16/bf16 shrink stored sizes further.
- Gradient accumulation— smaller forward peak, then accumulate (Ch 08-01).
torch.compile— Inductor can fuse checkpointed segments; compile whole model:
compiled = torch.compile(model_ckpt) # works, verify memory gains after profiling
-
Measuring Before/After¶
torch.cuda.reset_peak_memory_stats()
out.mean().backward()
torch.cuda.synchronize()
print("peak MiB:", torch.cuda.max_memory_allocated() // 2**20)
Run the same block with/without checkpointing; the peak delta is your real saving.
Pitfalls¶
- Recursion depth— heavily nested checkpointed blocks can use huge python stack; prefer few deep segments.
- Input mutation— checkpointed regions must not mutate their inputs in-place (recompute won't match).
- RNG—
checkpointuses a seed-fix to keep dropout deterministic: passpreserve_rng_statecarefully; nondeterministic layers (dropout with different masks across recompute) needuse_reentrant=False+ fixed seed. - Non-determinism in backward— if you observe gradient nondeterminism, check recompute of in-place ops.
Key Takeaways¶
- Checkpointing trades ~1 extra forward per segment for O(√N) activation memory.
- Wrap per transformer block;
use_reentrant=Falseis the modern safe flag. - Stack with FSDP + AMP + accumulation for the full memory toolkit.
- Measure peak with
max_memory_allocatedbefore trusting a config. - Keep inputs immutable inside checkpointed regions; mind RNG.
-