How do distributed training jobs checkpoint and recover after failure?

Distributed training checkpointing preserves a consistent snapshot of training progress so a failed worker, interrupted node, or restarted job can continue from a known good step instead of starting over. For faithful recovery, save more than model weights: optimizer and scheduler state, progress counters, random-number-generator (RNG) state, and the relevant data-sampling position must be preserved or reconstructed.

Question 107 covered resource budgets, quotas and preemption from the scheduling perspective. This lesson focuses on the technical checkpoint format, coordination, durable publication and restoration.

Beginner-friendly numerical example

Suppose a delivery-ETA model trains across 4 GPUs for 10,000 optimizer steps. Each step takes 0.5 seconds on average, excluding saving overhead.

• At step 4,000, the workers successfully commit checkpoint ckpt-4000.
• At step 4,650, one worker fails before the next successful checkpoint.
• After coordinated restart, the job loads the last fully committed step 4,000.
• Lost progress = 4650 - 4000 = 650 steps.
• Recompute time ≈ 650 × 0.5 = 325 seconds, about 5 minutes 25 seconds, before restart overhead.

Never resume at 4,650 just because logs report reaching that step: there is no verified checkpoint for the model and optimizer state at 4,650.

If checkpointing occurs every 1,000 steps, maximum lost progress from a failure immediately before the next save is approximately 999 steps. At 0.5 seconds per step, that is about 8 minutes 20 seconds. Actual loss also depends on checkpoint timing and whether asynchronous writes have completed.

What is inside a resumable checkpoint?

Model parameters: Weights and buffers, including sharded tensors in FSDP/ZeRO-style training.

Optimizer state: Momentum, Adam running averages and other parameter-specific variables. Reloading weights without optimizer state often changes the next training updates.

Scheduler and scaler: Learning-rate scheduler position and mixed-precision gradient-scaler state, if used.

Training progress: Completed epoch, global optimizer step, accumulation position, sample count and any stage-specific state.

Data-loader/sampler state: Current shard, shuffled permutation/seed, batch offset or equivalent replay position. Otherwise a restarted run can repeat or skip examples.

RNG state: Python random, NumPy, torch CPU and the appropriate CUDA/device RNG states, preferably per rank. Capture these consistently with training step boundaries.

Run identity: Dataset snapshot, training configuration, software/container version, distributed topology, world size and model/optimizer checkpoint format.

Integrity metadata: Checkpoint ID, completion manifest, expected shard list, checksums and committed timestamp. A stored directory is not enough evidence that all rank shards are valid.

For gradient accumulation interrupted midstep, either preserve gradient accumulation state or checkpoint only at a clearly defined optimizer-step boundary and replay from that boundary. Saving partial gradients from some workers but not others is not a consistent recovery point.

Coordinated snapshots across workers

Imagine four ranks: 0, 1, 2 and 3. In sharded training, each rank may own different model or optimizer shards. A usable checkpoint must correspond to the same logical step across all ranks.

1. Finish the target optimizer step and reach a coordinated checkpoint boundary.
2. Construct state dictionaries with the supported distributed checkpointing APIs.
3. Have participating ranks write their shards to a new, unique checkpoint location.
4. Validate that the distributed save finished across the participating ranks.
5. Publish an atomic or transactional commit marker / manifest only when all required pieces are durable and verified.
6. Advance the latest_committed pointer to the new checkpoint.
7. Keep the previous committed checkpoint available until the new one has passed verification.

PyTorch Distributed Checkpoint (DCP) supports parallel saving/loading and resharding on load. DeepSpeed explicitly warns that save_checkpoint must be called by all participating processes, not just rank 0, or the operation can hang due to synchronization. <See the linked PyTorch and DeepSpeed references.>

Do not mix shards from different steps. An apparently complete directory containing rank-0 state from step 5,000 and rank-1 state from step 4,000 is not a valid checkpoint.

Atomic and durable writes

A safe publication protocol looks like:

checkpoint-5000.tmp/
   rank0 shard
   rank1 shard
   rank2 shard
   rank3 shard
   metadata

Verify writes and consistency
         ↓
Publish committed manifest for step 5000
         ↓
Update latest_committed pointer
         ↓
Eligible for restart

In a POSIX-compatible filesystem, an atomic rename can help publish a completed checkpoint within the same filesystem, subject to durability guarantees. In object stores, a multi-object upload is not an atomic directory rename; write immutable objects and publish a small commit/manifest object only after validation. Readers should resolve the committed manifest rather than scanning for the newest-looking prefix.

Design for power loss and worker failure during upload. If ckpt-5000 is incomplete, the loader must reject it and fall back to the last committed checkpoint, for example ckpt-4000.

Checksums detect corruption; they do not by themselves guarantee that shards came from one consistent optimizer step. Record both identities and integrity.

PyTorch Distributed Checkpoint example

A compact conceptual pattern using modern PyTorch DCP is:

PYTHON
import torch
import torch.distributed.checkpoint as dcp

# Prerequisites: torch.distributed process group initialized,
# one process per expected rank, model and optimizer created,
# and correct device/rank assignment.
state = {
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
}

# Every rank participates in a collective distributed save.
dcp.save(state_dict=state, checkpoint_id="/durable/ckpt-4000")

# On coordinated restart, initialize model and optimizer first.
restored = {
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
}
dcp.load(state_dict=restored, checkpoint_id="/durable/ckpt-4000")
model.load_state_dict(restored["model"])
optimizer.load_state_dict(restored["optimizer"])

Important limitations: This demonstrates the DCP entrypoints, not a ready-to-run FSDP or ZeRO checkpoint recipe. Correct sharded state dictionaries should use APIs appropriate to the installed torch version and model topology. This snippet omits scheduler, scaler, RNG, sampler, progress, error handling, shared storage configuration, and a separate committed-manifest protocol. Follow the official DCP tutorial for real distributed setup, especially if resharding to a different world size.

For a single-process experiment, ordinary torch.save with appropriate durability handling may be enough. For sharded large models, use DCP or the distributed framework's native save/load mechanism.

Recovery after a worker failure

1. Detect failure: Collective communication error, node termination, unhealthy rank or scheduler notification.
2. Stop or coordinate surviving workers: Do not let some ranks continue while others restore an older optimizer step.
3. Select checkpoint: Resolve latest_committed; verify metadata, shard checksums and format compatibility.
4. Reconstruct topology: Restart the required ranks/process group and allocate compatible devices.
5. Load state: Restore model, optimizer, scheduler and scaler along with RNG and data position.
6. Verify progress: Ensure all ranks agree on global step and training configuration.
7. Resume at a safe boundary: Continue from the next step after the last committed optimizer update.
8. Validate recovery: Compare short resumed runs with an uninterrupted baseline where determinism is expected; monitor for duplicate data and divergent metrics.

An elastic restart can have a different number of workers. PyTorch DCP can reshard checkpoint state on load, but resharding alone does not preserve identical data-order or optimizer semantics when global batch size, sharding or learning-rate policy changes. Specify supported topology changes and test them.

RNG-state restoration

Restoring weights is not the same as restoring the random stream. Example RNG capture pattern:

PYTHON
import random
import numpy as np
import torch

rng_state = {
    "python": random.getstate(),
    "numpy": np.random.get_state(),
    "torch_cpu": torch.get_rng_state(),
    "torch_cuda": torch.cuda.get_rng_state_all()
        if torch.cuda.is_available() else None,
}

# On compatible restart after loading the model state:
random.setstate(rng_state["python"])
np.random.set_state(rng_state["numpy"])
torch.set_rng_state(rng_state["torch_cpu"])
if rng_state["torch_cuda"] is not None:
    torch.cuda.set_rng_state_all(rng_state["torch_cuda"])

This is an example of what to capture, not a complete serialization or distributed restoration implementation. Capture rank-specific RNG state and restore it after the process topology and generators are initialized, avoiding unintended random draws before training resumes. Modern NumPy code may use independent Generator instances whose bit_generator.state must be handled separately. Sampler and data-loader workers also require coordinated seeds/positions.

Exact bitwise replay may still be impossible across differing GPUs, nondeterministic kernels or framework versions. Distinguish recovering useful training progress from identical numerical replay.

Choose a checkpoint frequency: numerical trade-off

Assume each full coordinated checkpoint adds 12 seconds of blocking save overhead, and failures are rare and roughly uniformly distributed within each interval.

With a checkpoint every 20 minutes:

• Save overhead ≈ 12 seconds per 20 minutes = 1% of elapsed training time.
• Expected lost work after a failure ≈ 10 minutes, if failure timing is uniform.

With a checkpoint every 5 minutes:

• Save overhead ≈ 12 / 300 × 100 = 4%.
• Expected lost work ≈ 2.5 minutes.

More frequent saves reduce lost progress but consume additional I/O, storage and compute. The truly optimal interval depends on failure rate, save time, restart cost, checkpoint size and training speed; use measured values rather than fixed universal defaults.

Asynchronous saving can overlap storage writes with training, reducing blocking time, but increases temporary CPU memory pressure and operational complexity. PyTorch's async DCP recipe recommends controlling in-flight saves, often allowing just one outstanding asynchronous checkpoint at a time, and waiting for its upload to complete before declaring success.

What to monitor in production

• Last committed checkpoint ID and global step—not simply last attempted save.
• Save duration, upload duration, total checkpoint bytes and validation failures.
• Time since latest valid checkpoint and estimated recoverable work at risk.
• Number of failed workers, coordinated restarts and recovery time.
• Resume success rate, checkpoint compatibility errors and corruption incidents.
• Sample/epoch replay anomalies and unexpected jumps in metrics after restart.
• Durable storage availability and capacity sufficient for incomplete plus prior valid snapshots.

Fault-injection tests should deliberately terminate one worker during training and another during a checkpoint write. Verify that the job restarts from the last complete checkpoint, does not mix shards, and can continue without silently discarding optimizer state.

Common mistakes

• Saving only model weights and claiming identical optimizer resumption.
• Writing each rank's state independently without a coordinated step or completion protocol.
• Loading the newest directory even when one rank's upload failed.
• Letting rank 0 alone call a collective save API.
• Keeping checkpoints only on an ephemeral worker disk.
• Advancing latest before the durable write completes.
• Forgetting RNG state, data-loader progress or gradient-accumulation boundaries.
• Assuming a changed world size implies unchanged sample order and numerical results.
• Creating overlapping async checkpoints until CPU memory is exhausted.
• Never testing actual restore paths until the first production incident.

In short: Make checkpoints complete, coordinated, immutable until committed, durable and restorable. Save model and optimizer state alongside progress, RNG and data position; restart all required workers consistently; reject partial writes; and tune checkpoint cadence against expected lost work and save overhead.