Checkpointing and Recovery › Code Implementation
Language AI HandbookAll chapters Checkpointing and Recovery › Code ImplementationOne section · Sections in this chapter6 / 9 ↓
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 build a practical checkpointing system that demonstrates the core concepts: saving a complete checkpoint, loading it for recovery, and measuring the overhead of synchronous versus simulated asynchronous checkpointing.
Setup and Imports
We will work with a medium-sized transformer model to make the storage and timing numbers realistic. The model is deliberately sized to produce measurable I/O times without requiring GPU hardware.
import torch
import torch.nn as nn
import torch.optim as optim
torch.manual_seed(42)
# Simulate a medium-sized model (like a small transformer block stack)
class SmallTransformer(nn.Module):
def __init__(self, vocab_size=32000, d_model=512, nhead=8, num_layers=6):
super().__init__()
self.embedding = nn.Embedding(vocab_size, d_model)
encoder_layer = nn.TransformerEncoderLayer(
d_model=d_model,
nhead=nhead,
dim_feedforward=d_model * 4,
dropout=0.1,
batch_first=True,
)
self.transformer = nn.TransformerEncoder(
encoder_layer, num_layers=num_layers
)
self.head = nn.Linear(d_model, vocab_size)
def forward(self, x):
x = self.embedding(x)
x = self.transformer(x)
return self.head(x)
model = SmallTransformer()
optimizer = optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0.01)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10000)
scaler = torch.amp.GradScaler("cpu")Model parameters: 51,714,304 Parameter storage (BF16 equivalent): 206.9 MB
This gives a concrete sense of the scale. A model of this size has parameters, optimizer state, and gradients that together span several hundred megabytes even for a "small" 6-layer transformer. Production models are hundreds of times larger.
Building a Complete Checkpoint
A complete checkpoint function captures all the state described earlier: parameters, optimizer, scheduler, scaler, step number, and data loader position.
import os
from pathlib import Path
def save_checkpoint(
checkpoint_dir, model, optimizer, scheduler, scaler, step, data_seed
):
"""
Save a complete training checkpoint.
data_seed: integer representing the RNG seed + step, used to reconstruct
which data examples correspond to this training step.
"""
checkpoint_dir = Path(checkpoint_dir)
checkpoint_dir.mkdir(parents=True, exist_ok=True)
checkpoint = {
"step": step,
"model_state_dict": model.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
"scheduler_state_dict": scheduler.state_dict(),
"scaler_state_dict": scaler.state_dict(),
"rng_state": torch.get_rng_state(),
"data_seed": data_seed,
"config": {
"d_model": 512,
"nhead": 8,
"num_layers": 6,
},
}
checkpoint_path = checkpoint_dir / f"checkpoint_step_{step:08d}.pt"
torch.save(checkpoint, checkpoint_path)
# Update the "latest" pointer atomically
latest_path = checkpoint_dir / "latest.txt"
tmp_path = checkpoint_dir / "latest.txt.tmp"
with open(tmp_path, "w") as f:
f.write(str(checkpoint_path))
os.replace(tmp_path, latest_path) # atomic on POSIX
return checkpoint_pathThe atomic update of latest.txt is important. If the process crashes while writing the checkpoint file itself, latest.txt still points to the previous valid checkpoint. Recovery logic always reads latest.txt first, so a partially written checkpoint file is never used. Without the atomic update (for example, if you updated latest.txt first and then wrote the checkpoint), a crash between those two operations would leave latest.txt pointing to a checkpoint that does not yet exist.
The os.replace() call is atomic on POSIX-compliant filesystems because it is implemented as a single rename system call. There is no window of time during which an observer could see latest.txt pointing to a nonexistent file. This guarantee holds for local filesystems and most distributed filesystems, but may not hold for all cloud object storage systems, which is why production systems sometimes use a dedicated coordination service (like etcd or Zookeeper) to manage the checkpoint pointer.
import os
import tempfile
import time
# Simulate a few training steps
train_data_seed = 12345
with tempfile.TemporaryDirectory() as tmpdir:
# Simulate training for a few steps
step = 1000
# Quick fake training step
dummy_input = torch.randint(0, 32000, (2, 64))
dummy_target = torch.randint(0, 32000, (2, 64))
optimizer.zero_grad()
output = model(dummy_input)
loss = nn.CrossEntropyLoss()(output.view(-1, 32000), dummy_target.view(-1))
loss.backward()
optimizer.step()
scheduler.step()
start = time.time()
ckpt_path = save_checkpoint(
tmpdir, model, optimizer, scheduler, scaler, step, train_data_seed
)
save_time = time.time() - start
ckpt_size_mb = os.path.getsize(ckpt_path) / 1e6
checkpoint_stats = {
"path": ckpt_path,
"size_mb": ckpt_size_mb,
"save_time_s": save_time,
"step": step,
}Checkpoint saved in 0.95s Checkpoint size: 620.7 MB Effective write speed: 651 MB/s Step: 1000
Loading a Checkpoint for Recovery
The recovery function loads the checkpoint and restores every piece of state. Notice how it mirrors the save function exactly, restoring each component in the reverse order it would be needed during training initialization.
def load_checkpoint(checkpoint_dir, model, optimizer, scheduler, scaler):
"""
Load the most recent valid checkpoint and restore all training state.
Returns the step to resume from, or 0 if no checkpoint exists.
"""
checkpoint_dir = Path(checkpoint_dir)
latest_path = checkpoint_dir / "latest.txt"
if not latest_path.exists():
print("No checkpoint found, starting from scratch.")
return 0, None
with open(latest_path, "r") as f:
ckpt_file = f.read().strip()
if not os.path.exists(ckpt_file):
print(f"Checkpoint file {ckpt_file} not found, starting from scratch.")
return 0, None
checkpoint = torch.load(ckpt_file, weights_only=True)
model.load_state_dict(checkpoint["model_state_dict"])
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
scaler.load_state_dict(checkpoint["scaler_state_dict"])
torch.set_rng_state(checkpoint["rng_state"])
step = checkpoint["step"]
data_seed = checkpoint["data_seed"]
return step, data_seed# Test the recovery
model_recovered = SmallTransformer()
optimizer_recovered = optim.AdamW(
model_recovered.parameters(), lr=1e-4, weight_decay=0.01
)
scheduler_recovered = optim.lr_scheduler.CosineAnnealingLR(
optimizer_recovered, T_max=10000
)
scaler_recovered = torch.amp.GradScaler("cpu")
with tempfile.TemporaryDirectory() as tmpdir:
# Save a checkpoint (re-run since tmpdir is new)
ckpt_path2 = save_checkpoint(
tmpdir, model, optimizer, scheduler, scaler, step, train_data_seed
)
start = time.time()
recovered_step, recovered_seed = load_checkpoint(
tmpdir,
model_recovered,
optimizer_recovered,
scheduler_recovered,
scaler_recovered,
)
load_time = time.time() - start
# Verify parameter equality
params_match = all(
torch.allclose(p1, p2)
for p1, p2 in zip(model.parameters(), model_recovered.parameters())
)
# Verify optimizer state equality
orig_lr = optimizer.param_groups[0]["lr"]
recov_lr = optimizer_recovered.param_groups[0]["lr"]
recovery_stats = {
"load_time_s": load_time,
"recovered_step": recovered_step,
"params_match": params_match,
"orig_lr": orig_lr,
"recov_lr": recov_lr,
"lr_match": abs(orig_lr - recov_lr) < 1e-10,
}Checkpoint loaded in 0.11s Resumed from step: 1000 Parameters match: True Learning rate match: True (lr=1.00e-04)
The params_match: True and lr_match: True outputs confirm that the recovery function restores the model to a state that is bit-for-bit identical to the original. This is the ground truth for a correct checkpoint implementation.
Measuring Checkpoint Overhead and Frequency Tradeoffs
Let's simulate the overhead tradeoff at different checkpoint frequencies and model sizes to understand when asynchronous checkpointing becomes critical.
# Simulate checkpoint timing for different model sizes
# Step time represents one training step (forward + backward + optimizer step)
step_time_ms = 800 # ~800ms per step for a large model on GPU
# Model sizes: number of parameters in billions
model_sizes_b = [1, 7, 13, 30, 70]
# Storage per parameter: BF16 (2 bytes) + Adam state (2 * 4 bytes) = 10 bytes
bytes_per_param = 10
# Assume 10 GB/s write bandwidth (NFS or cloud storage)
write_bandwidth_gbps = 10.0
checkpoint_data = []
for size_b in model_sizes_b:
size_bytes = size_b * 1e9 * bytes_per_param
size_gb = size_bytes / 1e9
write_time_s = size_gb / write_bandwidth_gbps
write_time_ms = write_time_s * 1000
for interval_steps in [50, 100, 250, 500, 1000]:
interval_time_ms = interval_steps * step_time_ms
overhead_pct = (write_time_ms / interval_time_ms) * 100
checkpoint_data.append(
{
"model_size_b": size_b,
"interval_steps": interval_steps,
"write_time_s": write_time_s,
"overhead_pct": overhead_pct,
}
) Model | Interval | Write Time | Overhead %
--------------------------------------------------
7B | 100 steps | 7.0s | 8.8%
7B | 500 steps | 7.0s | 1.8%
70B | 100 steps | 70.0s | 87.5%
70B | 500 steps | 70.0s | 17.5%The table makes the problem concrete. A 70B model writing to 10 GB/s storage requires 70 seconds per checkpoint. At a 100-step interval (about 80 seconds of training at 800ms per step), the synchronous overhead exceeds 87%. Asynchronous checkpointing is not a nice-to-have in this regime; it is a requirement for practical operation.
Visualizing Checkpoint Frequency vs. Overhead
Visualizing Expected Work Lost vs. Checkpoint Interval
The two charts together tell a complete story. The overhead chart shows that synchronous checkpointing at short intervals is infeasible for large models. The wasted-work chart shows that long intervals cause unacceptable work loss at scale. The only way to satisfy both constraints simultaneously is asynchronous checkpointing, which eliminates most of the overhead and allows short intervals without the associated I/O cost.
Rolling Checkpoint Policy
A rolling checkpoint manager retains only the last checkpoints, deleting older ones automatically. The implementation also maintains permanent checkpoints at specified intervals, giving long-term fallback protection.
class RollingCheckpointManager:
"""
Manages a rolling window of checkpoints, retaining only the last K.
Also maintains permanent checkpoints at specified intervals.
"""
def __init__(self, checkpoint_dir, keep_last=5, permanent_every=1000):
self.checkpoint_dir = Path(checkpoint_dir)
self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
self.keep_last = keep_last
self.permanent_every = permanent_every
self.history = [] # list of (step, path) tuples
def save(self, model, optimizer, scheduler, scaler, step, data_seed):
is_permanent = step % self.permanent_every == 0
prefix = "permanent" if is_permanent else "rolling"
checkpoint = {
"step": step,
"model_state_dict": model.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
"scheduler_state_dict": scheduler.state_dict(),
"scaler_state_dict": scaler.state_dict(),
"rng_state": torch.get_rng_state(),
"data_seed": data_seed,
"is_permanent": is_permanent,
}
path = self.checkpoint_dir / f"{prefix}_step_{step:08d}.pt"
torch.save(checkpoint, path)
if not is_permanent:
self.history.append((step, path))
# Delete old rolling checkpoints beyond keep_last
while len(self.history) > self.keep_last:
old_step, old_path = self.history.pop(0)
if old_path.exists():
old_path.unlink()
# Update latest pointer
latest_path = self.checkpoint_dir / "latest.txt"
tmp_path = self.checkpoint_dir / "latest.txt.tmp"
with open(tmp_path, "w") as f:
f.write(str(path))
os.replace(tmp_path, latest_path)
return path, is_permanent# Demonstrate rolling policy behavior
with tempfile.TemporaryDirectory() as tmpdir:
manager = RollingCheckpointManager(tmpdir, keep_last=3, permanent_every=500)
saved_records = []
for step_i in [100, 200, 300, 400, 500, 600, 700]:
path, is_perm = manager.save(
model, optimizer, scheduler, scaler, step_i, 42
)
exists = path.exists()
saved_records.append(
{
"step": step_i,
"permanent": is_perm,
"exists": exists,
}
)
# Check what's still on disk after all saves
existing_files = sorted(Path(tmpdir).glob("*.pt"))
for record in saved_records:
# Re-check existence (rolling may have deleted)
record["exists_after"] = any(
f"step_{record['step']:08d}" in str(p) for p in existing_files
)Step | Permanent | Retained ---------------------------------- 100 | no | deleted 200 | no | deleted 300 | no | deleted 400 | no | yes 500 | yes | yes 600 | no | yes 700 | no | yes Files still on disk: 4 permanent_step_00000500.pt rolling_step_00000400.pt rolling_step_00000600.pt rolling_step_00000700.pt
The output shows the rolling policy in action: steps 100 and 200 were deleted once step 700 was saved (exceeding the keep_last=3 window), but step 500 was retained permanently. A team investigating a regression discovered six days later can roll back to the step-500 permanent checkpoint even after hundreds of subsequent rolling checkpoints have been purged. The two-tier policy provides both operational recovery (rolling checkpoints for recent failures) and research-grade recovery (permanent checkpoints for long-term fallback).
Checkpoint Timing Breakdown
Let's measure the actual time breakdown between the CPU copy phase and the disk write phase to understand the async checkpointing opportunity precisely.
import io
def measure_checkpoint_phases(model, optimizer):
"""
Measure time for each phase of checkpointing to understand async opportunity.
"""
results = {}
# Phase 1: Serialize to in-memory buffer (simulates CPU copy from GPU)
buf = io.BytesIO()
t0 = time.perf_counter()
torch.save(
{
"model": model.state_dict(),
"optimizer": optimizer.state_dict(),
},
buf,
)
t1 = time.perf_counter()
results["serialize_ms"] = (t1 - t0) * 1000
results["size_mb"] = buf.tell() / 1e6
# Phase 2: Write from memory buffer to disk (simulates background I/O thread)
buf.seek(0)
with tempfile.NamedTemporaryFile(delete=True, suffix=".pt") as f:
t2 = time.perf_counter()
f.write(buf.read())
f.flush()
os.fsync(f.fileno())
t3 = time.perf_counter()
results["write_ms"] = (t3 - t2) * 1000
results["total_sync_ms"] = results["serialize_ms"] + results["write_ms"]
results["async_stall_ms"] = results[
"serialize_ms"
] # only stall during serialize
results["async_savings_pct"] = (
results["write_ms"] / results["total_sync_ms"]
) * 100
return resultstiming_results = measure_checkpoint_phases(model, optimizer)Checkpoint Timing Breakdown: Serialize (CPU copy equivalent): 200.7ms Write to disk: 208.6ms Total (synchronous): 409.3ms Async stall (serialize only): 200.7ms Time savings from async I/O: 51% Checkpoint size: 620.7MB
The timing breakdown confirms the core intuition: the disk write phase dominates total checkpoint time. Asynchronous checkpointing eliminates this dominant cost from the training stall, reducing the observable pause to just the serialization time. Even on a local SSD, disk writes take substantially longer than in-memory serialization for the same data, and the gap grows considerably on network-attached storage or cloud object stores.
Synchronous vs. Asynchronous Stall Comparison
To make the benefit of async checkpointing concrete across model scales, let's visualize how the synchronous stall (full serialize plus write) compares to the asynchronous stall (serialize only) for models ranging from 1B to 70B parameters.
The reduction for the 70B model is substantial. At that scale, the disk write alone consumes 70 seconds, while the serialization takes only about 11 seconds. Asynchronous checkpointing recovers those 70 seconds of idle GPU time entirely, at the cost of holding an extra copy of the model state in CPU memory.