Checkpointing & Fault Tolerance¶
Overview¶
Distributed training is a marathon: hours-to-weeks of steps where a single GPU OOM, node loss, or power event can wipe progress. Checkpointing is your insurance policy— done right it costs little and saves days. The rules differ between plain, DDP, and FSDP.
- Save model + optimizer + RNG + epoch/batch + sampler state— nothing less.
- Only one rank writes (or all ranks write identical files) to avoid races.
- Load must be rank-consistent: every rank restores the same logical state.
- Atomic saves (tmp file + rename) prevent corrupt checkpoints after crashes.
The most expensive bug in distributed training: checkpointing "the model" but not the optimizer/sampler → resumed training diverges from expected behavior.
-
What to Save (the full recipe)¶
import torch
def build_checkpoint(model, optimizer, epoch, step, sampler_state, extra=None):
return {
"model": model.state_dict(),
"optimizer": optimizer.state_dict(),
"epoch": epoch,
"step": step,
"sampler": sampler_state, # DistributedSampler state!
"rng": {
"torch": torch.get_rng_state(),
"cuda": torch.cuda.get_rng_state_all(),
},
"extra": extra or {},
}
Without
sampler.set_epoch(epoch)+ sampler state, resuming duplicates/skips data.
Saving Correctly (plain & DDP)¶
Single-rank save (simplest, works everywhere)¶
if dist.get_rank() == 0: # only leader writes
tmp = f"ckpt_{step}.pt.tmp"
torch.save(build_checkpoint(...), tmp)
os.replace(tmp, f"ckpt_{step}.pt") # atomic
All-rank save with identical content (DDP-safe)¶
state = build_checkpoint(model, optimizer, epoch, step, sampler_state)
torch.save(state, f"ckpt_rank{dist.get_rank()}.pt") # same logical state, per-rank file
# or
Loading Correctly¶
ckpt = torch.load("ckpt.pt", map_location="cpu") # always load to CPU first!
model.load_state_dict(ckpt["model"])
optimizer.load_state_dict(ckpt["optimizer"])
epoch, step = ckpt["epoch"], ckpt["step"]
torch.set_rng_state(ckpt["rng"]["torch"])
torch.cuda.set_rng_state_all(ckpt["rng"]["cuda"])
sampler.set_epoch(epoch) # resume data order
Load to CPU then
.to(device)— avoids device-mismatch errors and huge host copies.
FSDP Checkpointing— the Sharded Contract¶
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
# Save (all ranks participate; each writes its shard):
fsdp_state = {
"model": model.state_dict(), # sharded
"optim": optimizer.state_dict(), # sharded
}
if rank == 0:
torch.save(fsdp_state, "fsdp.pt") # single file, model handles sharding
# Load (all ranks):
ckpt = torch.load("fsdp.pt")
model.load_state_dict(ckpt["model"]) # auto-reshard onto this rank
optimizer.load_state_dict(ckpt["optim"])
Use the same FSDP wrapping structure (same auto_wrap_policy, same world_size at load)— or use
ShardedStateDict+reshardAPIs for flexible world sizes.
Asynchronous & Continuous Checkpointing¶
Saving a 100GB model blocks training for minutes— decouple it:
# Dump state dict to CPU (cheap), save on a background thread:
state = build_checkpoint(...) # on-GPU tensors already CPU in state_dict
import threading
t = threading.Thread(target=lambda: torch.save(state, f"ckpt_{step}.pt"))
t.start() # training continues immediately
Production systems (NeMo, torchtitan) do continuous async saves every N steps, keeping N-1 checkpoints to guard against corruption.
Fault Tolerance Patterns¶
| Failure | Mitigation |
|---|---|
| GPU OOM mid-training | smaller batch, activation checkpointing (Ch 08-03), gradient accumulation |
| Node lost (SLURM/k8s) | checkpoint every N steps; restart from last good |
| Corrupt checkpoint | atomic rename + write to tmp dir + checksums |
| Rank hang (deadlock) | timeouts (init_process_group(..., timeout=...)), watchdog |
Detecting divergence before it costs you¶
def check_loss_sanity(loss, rank, threshold=1e6):
if rank == 0 and not torch.isfinite(loss):
print(f"WARNING: non-finite loss at step {step}")
# optionally: save emergency checkpoint before optimizer corrupts further
-
Validation & Restart Drill¶
- Save a checkpoint at step 100.
- Kill training (Ctrl-C / kill -9 simulation).
- Restart from that checkpoint, run 10 more steps.
- Re-run steps 100-110 without checkpointing from step 99 version.
- Assert losses match within float tolerance:
assert torch.allclose(loss_a, loss_b, atol=1e-4), "checkpoint resume is broken"
This drill catches 90% of checkpoint bugs before a 3-day run eats them.
-
Key Takeaways¶
- Save model + optimizer + RNG + epoch/step + sampler state and nothing less.
- Leader-rank saves only; atomic writes; load to CPU first.
- FSDP: save/load
state_dict()on all ranks; keep wrap structure identical. - Async/continuous saving decouples I/O from training.
- Run a small restart drill to prove resume correctness before long runs.
-