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: checkpoint the big transformer block (attention + mlp), recompute inside
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)

1. FSDP — FSDP shards params; checkpointing shards activations: the standard LLM combo.
2. AMP — activations in fp16/bf16 shrink stored sizes further.
3. Gradient accumulation — smaller forward peak, then accumulate (Ch 08-01).
4. 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

1. Recursion depth — heavily nested checkpointed blocks can use huge python stack; prefer few deep segments.
2. Input mutation — checkpointed regions must not mutate their inputs in-place (recompute won't match).
3. RNG — checkpoint uses a seed-fix to keep dropout deterministic: pass preserve_rng_state carefully; nondeterministic layers (dropout with different masks across recompute) need use_reentrant=False + fixed seed.
4. 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=False is the modern safe flag.
• Stack with FSDP + AMP + accumulation for the full memory toolkit.
• Measure peak with max_memory_allocated before trusting a config.
• Keep inputs immutable inside checkpointed regions; mind RNG.

---

Related Topics

• Gradient Accumulation
• Fsdp
• Amp