Skip to content

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: collect on rank 0 -> single file

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 + reshard APIs 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

  1. Save a checkpoint at step 100.
  2. Kill training (Ctrl-C / kill -9 simulation).
  3. Restart from that checkpoint, run 10 more steps.
  4. Re-run steps 100-110 without checkpointing from step 99 version.
  5. 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.