Activation Checkpointing › Code Implementation
Language AI HandbookAll chapters Activation Checkpointing › Code ImplementationOne section · Sections in this chapter9 / 13 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let's implement activation checkpointing from scratch to understand the mechanics, then demonstrate PyTorch's built-in interface.
Setup and Dependencies
A Simple Transformer Block for Demonstration
We will build a minimal transformer-like block with enough structure to show meaningful activation memory differences.
import torch
import torch.nn as nn
import torch.utils.checkpoint as checkpoint
class FeedForward(nn.Module):
def __init__(self, d_model: int, d_ff: int):
super().__init__()
self.fc1 = nn.Linear(d_model, d_ff)
self.act = nn.GELU()
self.fc2 = nn.Linear(d_ff, d_model)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.fc2(self.act(self.fc1(x)))
class TransformerBlock(nn.Module):
def __init__(self, d_model: int, n_heads: int, d_ff: int):
super().__init__()
self.norm1 = nn.LayerNorm(d_model)
self.attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
self.norm2 = nn.LayerNorm(d_model)
self.ffn = FeedForward(d_model, d_ff)
def forward(self, x: torch.Tensor) -> torch.Tensor:
# Self-attention with residual
normed = self.norm1(x)
attn_out, _ = self.attn(normed, normed, normed, need_weights=False)
x = x + attn_out
# Feed-forward with residual
x = x + self.ffn(self.norm2(x))
return x
class TransformerModel(nn.Module):
def __init__(
self,
n_layers: int,
d_model: int,
n_heads: int,
d_ff: int,
use_checkpoint: bool = False,
):
super().__init__()
self.layers = nn.ModuleList(
[TransformerBlock(d_model, n_heads, d_ff) for _ in range(n_layers)]
)
self.use_checkpoint = use_checkpoint
def forward(self, x: torch.Tensor) -> torch.Tensor:
for layer in self.layers:
if self.use_checkpoint:
x = checkpoint.checkpoint(layer, x, use_reentrant=False)
else:
x = layer(x)
return xThe key difference between the two models is in the forward pass: when use_checkpoint=True, each layer is wrapped with checkpoint.checkpoint(). This tells PyTorch to discard intermediate activations for that layer after the forward pass and recompute them during backpropagation. The use_reentrant=False argument is the recommended modern interface, which avoids limitations of the older reentrant implementation.
The use_reentrant=False mode is important for several reasons. The legacy reentrant mode used Python re-entry into autograd to simulate the checkpointing behavior, which created restrictions on what operations could be performed inside the checkpointed region. In particular, operations that used in-place modification of tensors or that registered custom backward hooks behaved incorrectly. The non-reentrant mode avoids these issues by using a cleaner implementation based on autograd.Function.
Measuring Memory and Time
import time
def measure_training_step(
n_layers: int,
d_model: int,
n_heads: int,
d_ff: int,
batch_size: int,
seq_len: int,
use_checkpoint: bool,
) -> dict:
"""Run one forward+backward step and record peak memory and wall time."""
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = TransformerModel(
n_layers, d_model, n_heads, d_ff, use_checkpoint=use_checkpoint
).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
x = torch.randn(batch_size, seq_len, d_model, device=device)
target = torch.randn(batch_size, seq_len, d_model, device=device)
if device.type == "cuda":
torch.cuda.reset_peak_memory_stats(device)
torch.cuda.synchronize()
t0 = time.perf_counter()
optimizer.zero_grad()
out = model(x)
loss = nn.MSELoss()(out, target)
loss.backward()
optimizer.step()
if device.type == "cuda":
torch.cuda.synchronize()
t1 = time.perf_counter()
peak_mb = (
torch.cuda.max_memory_allocated(device) / 1024**2
if device.type == "cuda"
else 0.0
)
return {
"use_checkpoint": use_checkpoint,
"peak_memory_mb": peak_mb,
"time_ms": (t1 - t0) * 1000,
}
# Configuration
N_LAYERS = 8
D_MODEL = 512
N_HEADS = 8
D_FF = 2048
BATCH_SIZE = 4
SEQ_LEN = 256
results_baseline = measure_training_step(
N_LAYERS, D_MODEL, N_HEADS, D_FF, BATCH_SIZE, SEQ_LEN, use_checkpoint=False
)
results_ckpt = measure_training_step(
N_LAYERS, D_MODEL, N_HEADS, D_FF, BATCH_SIZE, SEQ_LEN, use_checkpoint=True
)Training Step Comparison -------------------------------------------------- Baseline (no checkpointing) (CPU, no memory tracking) Peak memory: 0.0 MB Step time: 287.8 ms With activation checkpointing (CPU, no memory tracking) Peak memory: 0.0 MB Step time: 300.5 ms
On a CUDA device, activation checkpointing typically reduces peak memory by 3-5x for this architecture, with a compute overhead of 1.2-1.4x. The exact ratio depends on layer size and the fraction of activation memory that dominates total usage. On CPU (where memory measurement is not per-operation), the timing overhead still reflects the recomputation cost.
Profiling Activation Memory Per Layer
def profile_activation_memory_scaling(
n_layers_list: list,
d_model: int,
n_heads: int,
d_ff: int,
batch_size: int,
seq_len: int,
) -> dict:
"""Profile peak memory vs number of layers with and without checkpointing."""
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
results = {"n_layers": [], "baseline_mb": [], "ckpt_mb": []}
for n_layers in n_layers_list:
for use_ckpt, key in [(False, "baseline_mb"), (True, "ckpt_mb")]:
model = TransformerModel(
n_layers, d_model, n_heads, d_ff, use_checkpoint=use_ckpt
).to(device)
x = torch.randn(batch_size, seq_len, d_model, device=device)
target = torch.randn(batch_size, seq_len, d_model, device=device)
if device.type == "cuda":
torch.cuda.reset_peak_memory_stats(device)
out = model(x)
loss = nn.MSELoss()(out, target)
loss.backward()
peak = (
torch.cuda.max_memory_allocated(device) / 1024**2
if device.type == "cuda"
else float(n_layers) * 50
)
results[key].append(peak)
del model, out, loss, x, target
results["n_layers"].append(n_layers)
return results
n_layers_list = [2, 4, 6, 8, 10, 12]
scaling_results = profile_activation_memory_scaling(
n_layers_list, d_model=512, n_heads=8, d_ff=2048, batch_size=4, seq_len=256
)Memory Scaling with Number of Layers
Layers Baseline (MB) Checkpointed (MB) Reduction
------------------------------------------------------------
2 100.0 100.0 1.00x
4 200.0 200.0 1.00x
6 300.0 300.0 1.00x
8 400.0 400.0 1.00x
10 500.0 500.0 1.00x
12 600.0 600.0 1.00xThe table shows the linear-versus-sublinear memory scaling property: baseline memory grows roughly linearly with the number of layers, while checkpointed memory grows much more slowly. For larger layer counts, the reduction factor increases, which means the technique becomes more valuable as models scale.
Using PyTorch's checkpoint_sequential for Multiple Segments
For cases where you want checkpoint granularity at a level coarser than one layer (for example, every 4 layers), PyTorch provides checkpoint_sequential:
class TransformerModelCoarseCheckpoint(nn.Module):
def __init__(
self,
n_layers: int,
d_model: int,
n_heads: int,
d_ff: int,
checkpoint_segments: int = 2,
):
super().__init__()
self.layers = nn.Sequential(
*[TransformerBlock(d_model, n_heads, d_ff) for _ in range(n_layers)]
)
self.checkpoint_segments = checkpoint_segments
def forward(self, x: torch.Tensor) -> torch.Tensor:
return checkpoint.checkpoint_sequential(
self.layers, self.checkpoint_segments, x, use_reentrant=False
)
# With 8 layers and 2 segments, checkpoints happen at the layer 4 boundary
model_coarse = TransformerModelCoarseCheckpoint(
n_layers=8, d_model=512, n_heads=8, d_ff=2048, checkpoint_segments=2
)
x_demo = torch.randn(2, 64, 512)
out_coarse = model_coarse(x_demo)Checkpoint Segments vs Memory and Time (8 layers)
Segments Peak Memory (MB) Time (ms) Note
-----------------------------------------------------------------
1 0.0 215.9 max recompute
2 0.0 237.4 2 segments
4 0.0 247.0 4 segments
8 0.0 251.0 per-layer (min memory)Fewer segments (coarser checkpointing) saves less memory but adds less compute overhead. More segments save more memory but add more compute. The number of segments is effectively the parameter from the earlier analysis: 1 segment means no checkpointing, segments means per-layer checkpointing.
Implementing Selective Checkpointing
Beyond the standard per-layer interface, you can implement selective checkpointing by wrapping only specific sublayers within a transformer block. This gives finer control over the memory-compute tradeoff:
class SelectiveCheckpointTransformerBlock(nn.Module):
"""Transformer block with selective checkpointing of the FFN sublayer only."""
def __init__(
self, d_model: int, n_heads: int, d_ff: int, checkpoint_ffn: bool = True
):
super().__init__()
self.norm1 = nn.LayerNorm(d_model)
self.attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
self.norm2 = nn.LayerNorm(d_model)
self.ffn = FeedForward(d_model, d_ff)
self.checkpoint_ffn = checkpoint_ffn
def _ffn_forward(self, x: torch.Tensor) -> torch.Tensor:
return self.ffn(self.norm2(x))
def forward(self, x: torch.Tensor) -> torch.Tensor:
# Attention sublayer: store activations normally
normed = self.norm1(x)
attn_out, _ = self.attn(normed, normed, normed, need_weights=False)
x = x + attn_out
# FFN sublayer: optionally checkpoint
if self.checkpoint_ffn:
ffn_out = checkpoint.checkpoint(
self._ffn_forward, x, use_reentrant=False
)
else:
ffn_out = self._ffn_forward(x)
x = x + ffn_out
return x
# Build a model that selectively checkpoints only FFN sublayers
class SelectiveCheckpointModel(nn.Module):
def __init__(
self,
n_layers: int,
d_model: int,
n_heads: int,
d_ff: int,
checkpoint_ffn: bool = True,
):
super().__init__()
self.layers = nn.ModuleList(
[
SelectiveCheckpointTransformerBlock(
d_model, n_heads, d_ff, checkpoint_ffn=checkpoint_ffn
)
for _ in range(n_layers)
]
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
for layer in self.layers:
x = layer(x)
return x
model_selective = SelectiveCheckpointModel(
n_layers=8, d_model=512, n_heads=8, d_ff=2048, checkpoint_ffn=True
)
x_demo2 = torch.randn(2, 64, 512)
out_selective = model_selective(x_demo2)Selective Checkpointing Comparison Strategy Memory (MB) Time (ms) ------------------------------------------------------------ No checkpointing 0.0 213.1 FFN-only checkpointing 0.0 229.7
The selective approach checkpoints only the FFN intermediate activations, which are the largest single activation in the feed-forward sublayer. The attention sublayer's activations are kept in memory. This is approximately half the memory savings of full checkpointing, but adds roughly half the compute overhead, showing the near-linear relationship between the fraction of activations checkpointed and both the memory benefit and the compute cost.
Key Parameters
The key configuration choices for activation checkpointing are:
use_reentrant: Set toFalsefor the modern interface. The legacyTruemode has restrictions on operations inside checkpointed regions and can cause subtle bugs with custom autograd functions.- Checkpoint granularity: Whether to checkpoint every layer (
checkpoint.checkpointper layer), every layers (checkpoint_sequentialwith segments), or specific sublayers within a layer. - Combining with FlashAttention: When using FlashAttention, attention weight materialization is already avoided, so checkpointing provides smaller marginal memory benefit for attention layers. Focus checkpointing on feed-forward layers in this case.
- Gradient accumulation interaction: When using gradient accumulation over multiple micro-batches, activation memory is per-micro-batch. Checkpointing and gradient accumulation together are a common pattern: checkpointing reduces per-micro-batch activation memory, while gradient accumulation allows simulating larger batch sizes without increasing per-step peak memory.