AAI Agentic Handbook
Question 06 · ML infrastructure & serving

Design checkpointing and preemption handling for large training jobs.

Reviewed Oct 2026
How to spend the time
  1. 5 min — failure rates at scale, goodput as the metric
  2. 10 min — what a checkpoint contains; sharded and async writes
  3. 10 min — interval maths and storage tiers
  4. 10 min — preemption, detection, automatic restart
  5. 5 min — validation, ownership
01Concept

At scale, failure is routine: Meta reported 466 job interruptions in a 54-day window of Llama 3 405B pretraining on 16k H100s, most from hardware (GPUs, HBM). The metric is goodput — useful compute over paid compute.

A checkpoint is model + optimizer + dataloader position + RNG state + LR-scheduler step. Techniques: sharded checkpoints (each rank writes its shard in parallel); async checkpointing (snapshot to host memory, upload in the background — PyTorch DCP async_save); in-memory replicas on peers for fast recovery; tiered storage (local NVMe → parallel file system → object storage). The interval follows Young/Daly: T ≈ √(2 × C × MTBF), where C is the blocking cost of one checkpoint.

02Practical use case

70B on 1,024 H100s: a full checkpoint is ~1 TB (140 GB BF16 weights + ~840 GB FP32 master and Adam state). Async snapshot stalls training ~20–30 s; the upload finishes in 2–3 min. Interval 45–60 min; keep the last three plus daily for 14 days. On preemption (AWS spot gives a 2-minute notice, GCP spot ~30 s), the SIGTERM handler writes an emergency checkpoint if it fits. Restart is automatic: health checks evict bad nodes, 2–3% hot spares replace them, and the job resumes from the latest checkpoint.

03Staff framing
  • Measure goodput, time to detect a failure (NCCL watchdogs default to minutes — tune them down), and time to restart (target < 10 min).
  • Validate checkpoints with periodic test loads and checksums. A corrupt checkpoint discovered three days later is the worst outcome.
  • Failure modes: silent data corruption on a bad GPU (loss spikes or NaNs — per-rank checks); checkpoint storms when many jobs save in the same minute (add jitter); dataloader state not saved, so data repeats after restart.
  • Elastic resume with a different data-parallel size needs a resharding-capable format such as DCP.
  • Org scope: the training-infra team owns orchestration and defaults; buy storage (Lustre, WEKA, FSx, object storage) and build the restart automation.
04Estimate

Assumptions: 70B on 1,024 H100s; failure rate scaled from Llama 3 (466 interruptions / 54 days / 16,384 GPUs) → 466 × (1,024 / 16,384) / 54 ≈ 0.54/day → MTBF ≈ 44 h; checkpoint ≈ 1 TB; async stall C ≈ 30 s; restart R ≈ 10 min; illustrative $2.50/H100-hour.

  • Interval: √(2 × 30 s × 158,400 s) ≈ 3,080 s ≈ 51 min.
  • Waste: C/T + (T/2 + R)/MTBF = 0.5/51 + (25.5 + 10)/2,640 ≈ 1.0% + 1.3% ≈ 2.3%.
  • $: 1,024 × 730 × $2.50 ≈ $1.87M/month × 2.3% ≈ $43k/month lost. With a synchronous 5-min stall the optimum moves to ~2.7 h and waste to ~6.5% ≈ $120k/month — the case for async checkpointing in one line.
  • Storage: 1 TB × (3 rolling + 14 daily) ≈ 17 TB; writing 1 TB in 2–3 min needs ~6–8 GB/s sustained.
03060901201800246810Young/Daly optimumCheckpoint interval (min) →Cluster time wasted (%) →
Too frequent wastes stall time; too rare wastes recomputation after each failure.

What the answer sounds like at each level

Mid

Saves a checkpoint every N steps and restarts from it on failure.

Senior

Uses sharded async checkpoints, saves full training state, picks an interval and handles preemption notices.

Staff

Derives the interval from MTBF and checkpoint cost, prices goodput loss, designs restart automation with hot spares, and validates checkpoints as a reliability practice.

Follow-ups the interviewer will push

1Is spot capacity at ~60% off worth it?

For jobs under ~256 GPUs with fast async checkpoints, usually yes. For large gangs, any single preempted node stops the whole job, so effective MTBF shrinks with node count; keep big runs on reserved capacity.

2Restart takes 40 minutes.

Profile it: rescheduling, container pull, checkpoint load (1 TB at 10 GB/s ≈ 100 s), NCCL init on 1k ranks, dataloader fast-forward. Fixes: hot spares, pre-pulled images, local copy of the latest checkpoint, peer-memory restore.

3Loss spikes right after a restore.

Check the RNG, dataloader and optimizer state were all restored, the LR scheduler step matches, and no shards came from different checkpoint versions.

Red flags interviewers listen for

  • Saves weights only, not optimizer and dataloader state.
  • Synchronous checkpoints with no interval reasoning.
  • Never test-loads checkpoints.
  • No automation — a human restarts the job.

How it's scored

DimensionWhat strong looks like
CompletenessFull training state, sharded and async, validated by test loads.
MathsYoung/Daly interval with real MTBF and stall numbers; goodput in $.
RecoveryFast detection, hot spares, automated restart, preemption handlers.
OperationsStorage tiers, retention, jitter, elastic resume.

Sources

  1. The Llama 3 Herd of Models — Llama Team, Meta, 2024paper
  2. Distributed Checkpoint (torch.distributed.checkpoint) — PyTorch, 2025docs