AppliedAIPrep logoAppliedAI/Prep
MLOps & ML Engineering / 05

Design a large-scale training pipeline that resumes cleanly after a node failure.

At thousand-GPU scale, hardware failure is the norm, not the exception, and a run that cannot resume wastes weeks. The signal is checkpointing strategy, deterministic resume, and minimizing lost work. Here is the fault-tolerant design.

Updated Aug 2026 · Grounded in real Applied AI Engineer interview loops and written to a senior-engineer editorial bar.

TL;DR: Assume failures are routine. Checkpoint model + optimizer + scheduler + data-loader position + RNG state periodically and asynchronously, so training barely pauses, to durable storage. On failure the scheduler restarts the job, every rank loads the latest consistent checkpoint, and training resumes from that step with the same data order. Tune checkpoint frequency to trade IO overhead against lost work, and detect silent failures (degraded GPU, NCCL hang), not just crashes.

How to approach it. State the premise first: with thousands of GPUs running for weeks, a node will fail, so the question is not "if" but "how little work do you lose and how do you guarantee resume is correct." Center the answer on what must live in a checkpoint to resume exactly (not just weights) and on minimizing both checkpoint overhead and re-done work.

A strong answer. What a correct checkpoint contains. Not just model weights. You also need optimizer state (Adam momentum and variance, or the loss spikes on resume), the learning-rate scheduler step, the data-loader position (so you neither re-train on the same samples nor skip some), and RNG state (for reproducible dropout and augmentation). Miss any one and resume is subtly wrong even though the job "comes back up."

Strategy.

  • Periodic, asynchronous checkpoints. Snapshot every N steps to durable, redundant storage (object store or parallel filesystem). Write in the background so the GPUs barely stall; for sharded training (FSDP or ZeRO) each rank writes its own shard and you assemble a consistent global checkpoint.
  • Resume path. A failure triggers the scheduler (Kubernetes or Slurm with elastic training, e.g. torchrun elastic) to restart the job. All ranks load the latest complete checkpoint and continue from that step. Guard against half-written files: write to a temp path, atomic rename, and retain the last K good checkpoints.
  • Frequency tradeoff. Frequent checkpoints cost IO and stall time but lose almost nothing on failure; infrequent ones are cheap but throw away more progress. Tune to the observed failure rate and the checkpoint cost so a single failure loses at most one interval.
  • Detect silent failures. Crashes are the easy case. The dangerous ones are silent: a degraded GPU, a NCCL collective that hangs, or corrupted gradients that poison training without crashing. Monitor gradient norms, per-step throughput, and collective-op timeouts; health-check and evict bad nodes before they spread.

Reproducibility falls out of the same machinery: versioned data, pinned seeds, and full state in the checkpoint mean a run can be resumed or rebuilt and a regression traced to a step.

rendering diagram…

Key takeaways

  • A correct checkpoint is weights plus optimizer, scheduler, data-loader position, and RNG state; weights alone resume wrong.
  • Asynchronous writes plus atomic rename keep throughput high and prevent loading a half-written file.
  • Checkpoint frequency is a tunable trade between IO cost and lost work, sized to the failure rate.
  • Silent failures (degraded GPU, NCCL hang, poisoned gradients) waste more time than crashes, so monitor for them explicitly.

What interviewers probe next.

  • "Why checkpoint optimizer state, not just weights?" Adam's momentum and variance are part of the trajectory; resuming with them reset causes a loss spike or divergence.
  • "How do you checkpoint a sharded (FSDP) model efficiently?" Each rank saves its parameter and optimizer shard in parallel; sharded checkpoints avoid gathering the full model onto one node.
  • "Async checkpoint risks?" You must snapshot a coherent step and never load a partial file, so use atomic rename and retain the last few good copies.
  • "Elastic training?" The job can resume on a different node count; frameworks re-rendezvous so a failed node is replaced and training continues.

Common mistakes.

  • Saving only weights, so resume loses optimizer, scheduler, and data-order state and diverges.
  • Synchronous checkpoints that stall every rank, killing throughput.
  • No protection against loading a half-written checkpoint after a crash mid-save.
  • Handling only hard crashes and missing silent failures (degraded GPU, NCCL hang) that quietly waste days.
HOW DID IT GO?
0
UP NEXT ON YOUR JOURNEY
DISCUSSION · 0

No comments yet — be the first to share your approach.