IntelliSensei | Fault-Tolerant Training Runs: Async Checkpointing with torch.distributed.checkpoint

A fine-tuning job that takes four hours on one GPU is a script. The same job spread over 64 GPUs for nine days is an operation, and operations fail: a node gets preempted, an NCCL collective times out, ECC errors take a card offline, a spot instance disappears mid-step. At realistic hardware failure rates, a multi-day run on a few hundred GPUs will be interrupted more than once. The question is never whether it breaks, but how many GPU-hours you lose each time it does.

Most teams we are called in to help have checkpointing that technically works and economically does not: a single rank gathers the full model to CPU, writes one enormous file, and every GPU in the cluster sits idle for four minutes while it happens. They then save infrequently because saving is expensive, so each failure costs hours of lost progress. This tutorial fixes both halves of that problem with torch.distributed.checkpoint (DCP), asynchronous saves, and an elastic restart path.

The economics: how often should you checkpoint?

There is a clean way to reason about the interval. If a checkpoint costs C seconds of stalled training and the run fails on average every M seconds, then with checkpoint interval T you lose roughly T/2 seconds of recomputed work per failure plus C per checkpoint. Total overhead per unit time is approximately:

overhead(T) = C/T + T/(2M)

Minimized at T = sqrt(2 * M * C) — the classic Young/Daly result. Plug in real numbers: a cluster with a mean time between failures of 8 hours (M = 28800) and a 90-second blocking checkpoint gives T ≈ 2280s, about 38 minutes, with roughly 8% total overhead.

Now make the checkpoint asynchronous so the blocking portion is 3 seconds instead of 90. The optimum drops to T ≈ 415s, and overhead falls to about 1.5%. That is the whole argument for the rest of this article: cheap checkpoints let you checkpoint often, and checkpointing often is what makes failures boring.

Measure your own two numbers rather than trusting these. C is the wall-clock stall you can see in step times around a save; M you get from your job scheduler's history.

Why `torch.save(model.state_dict())` breaks at scale

On a single GPU it is fine. On a sharded run (FSDP2, HSDP, tensor parallel) it has three specific failure modes:

Gathering is expensive and spiky. Pulling a full 70B bf16 model to rank 0 means moving ~140 GB across the network into one host's memory. It is slow, and it is the most common cause of OOM in an otherwise healthy job.
The checkpoint is welded to the topology. A monolithic file saved from 64 ranks is awkward to resume on 48 ranks after a partial node failure — and if you saved sharded files naively, it is impossible.
It saves the model but not the run. Optimizer state (which for AdamW is twice the model size), LR scheduler position, gradient scaler state, RNG state, dataloader position, and the step counter all have to come back too, or your "resume" is a silent restart of the data order.

DCP addresses all three: every rank writes its own shards in parallel to a checkpoint directory, the format carries enough metadata to reshard on load, and the API takes an arbitrary state-dict of stateful objects.

A correct DCP save and load

The modern pattern uses the Stateful protocol so that model and optimizer state-dicts are collected and set with the right distributed semantics automatically.

import torch
import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.stateful import Stateful
from torch.distributed.checkpoint.state_dict import (
    get_state_dict, set_state_dict, StateDictOptions,
)

class AppState(Stateful):
    """Everything needed to resume the run, not just the weights."""

    def __init__(self, model, optimizer, scheduler, step=0, epoch=0):
        self.model = model
        self.optimizer = optimizer
        self.scheduler = scheduler
        self.step = step
        self.epoch = epoch

    def state_dict(self):
        msd, osd = get_state_dict(
            self.model, self.optimizer,
            options=StateDictOptions(cpu_offload=True),
        )
        return {
            "model": msd,
            "optim": osd,
            "sched": self.scheduler.state_dict(),
            "step": self.step,
            "epoch": self.epoch,
            "rng": torch.cuda.get_rng_state(),
        }

    def load_state_dict(self, sd):
        set_state_dict(
            self.model, self.optimizer,
            model_state_dict=sd["model"],
            optim_state_dict=sd["optim"],
        )
        self.scheduler.load_state_dict(sd["sched"])
        self.step = sd["step"]
        self.epoch = sd["epoch"]
        torch.cuda.set_rng_state(sd["rng"].cpu())

Saving and loading are then symmetric, and every rank participates:

def save(app_state, path):
    dcp.save({"app": app_state}, checkpoint_id=path)

def load(app_state, path):
    # in-place: the existing sharded tensors define the target layout
    dcp.load({"app": app_state}, checkpoint_id=path)

Two details that matter more than they look:

dcp.load is in-place. You must construct the model and optimizer first, wrapped exactly as they will be used (same FSDP2/TP wrapping), then load into them. DCP reads the target's sharding from the live tensors and pulls only the slices each rank needs. That is what makes resharding across a different world size work.
The optimizer must have state before you load into it. A freshly constructed AdamW has empty exp_avg buffers, so there is nothing to load into. get_state_dict/set_state_dict handle this correctly; hand-rolled optimizer.state_dict() round-trips usually do not.

Async saves: hide the cost

dcp.async_save copies the state-dict to staging memory (pinned CPU by default), returns a future, and lets a background thread do the actual write while training continues. The only synchronous cost is the staging copy.

from concurrent.futures import Future

_pending: Future | None = None

def checkpoint_async(app_state, path):
    global _pending
    if _pending is not None:
        _pending.result()          # never overlap two saves
    _pending = dcp.async_save({"app": app_state}, checkpoint_id=path)
    return _pending

# in the training loop
for step, batch in enumerate(loader, start=start_step):
    loss = train_step(model, batch)
    app_state.step = step
    if step % CKPT_EVERY == 0:
        checkpoint_async(app_state, f"{CKPT_DIR}/step-{step}")

# before exiting
if _pending is not None:
    _pending.result()

Rules for using this safely:

Never let two async saves overlap. Block on the previous future first, as above. If you regularly find yourself waiting, your interval is shorter than your write bandwidth supports.
Do not mutate parameters while staging is in flight — that is the point of the staging copy, so let async_save finish its copy before the next optimizer step. The default pinned-memory staging handles this; custom StorageWriter setups sometimes do not.
Budget host RAM. Staging holds a CPU copy of the sharded state on each rank. For a 70B model with AdamW across 64 ranks that is a few GB per host — fine — but on a fat single node with few ranks it is not free.
Write to fast shared storage. Per-rank parallel writes will saturate a slow NFS mount. Object storage or a parallel filesystem, with one directory per checkpoint, is the expected shape.

Resharding: resume on a different world size

This is the feature that turns a failure into a shrug. Node 7 died, you have 56 GPUs instead of 64, and the queue will not give you a replacement for two hours. Because DCP checkpoints store logical tensor slices with metadata rather than rank-indexed blobs, you can build the model under the new mesh and load the same directory:

# original
torchrun --nnodes 8 --nproc_per_node 8 train.py --resume auto
# after losing a node
torchrun --nnodes 7 --nproc_per_node 8 train.py --resume auto

Nothing in the load path changes. The constraint is that the logical model must be identical — same architecture, same parameter names, same dtypes. Sharding strategy and world size may differ; the model may not.

The same mechanism gives you the offline conversion you will eventually want for handing a model to a serving team:

from torch.distributed.checkpoint.format_utils import dcp_to_torch_save

dcp_to_torch_save("ckpt/step-9000", "model-step-9000.pt")

Do that as a separate CPU job, not inside the training loop.

Elastic restarts and atomicity

Automatic recovery has two halves: the launcher must restart workers, and your script must find the right checkpoint.

torchrun \
  --nnodes 6:8 \
  --nproc_per_node 8 \
  --max-restarts 5 \
  --rdzv-backend c10d \
  --rdzv-endpoint $HEAD_NODE:29500 \
  train.py --resume auto

With an --nnodes range, torchrun re-forms the rendezvous with whatever nodes are alive and restarts every worker from scratch — which means your process must be able to reconstruct itself from a checkpoint on startup, every time. Write --resume auto as: list the checkpoint directory, take the highest step that has a completion marker, load it, and set the starting step.

And insist on that completion marker. A checkpoint interrupted mid-write is the worst failure mode there is, because it looks valid. Write to step-9000.tmp, barrier across all ranks, have rank 0 rename to step-9000 and drop a _COMPLETE file, then barrier again. Your resume logic ignores any directory without _COMPLETE. Keep the last two or three complete checkpoints plus a sparse archive (every Nth) so you can also roll back past a loss spike rather than only past a crash.

What people forget to save

Weights and optimizer are the easy part. Restarts that produce a subtle quality regression almost always trace to one of these:

Dataloader position. Restarting the epoch means re-training on data the model already saw and never reaching the tail of the shuffle. torchdata's StatefulDataLoader is a drop-in that adds state_dict()/load_state_dict(); without it, at minimum record the sample index and skip forward deterministically.
RNG state, per rank, for CUDA and CPU and the Python random module. Dropout masks and augmentation should continue, not restart.
LR scheduler step. Loading weights but restarting the warmup is a classic way to put a visible bump in your loss curve.
GradScaler state, if you are on fp16 AMP.
The token/sample counter you use for LR and data-mixing schedules, which may not equal the optimizer step under gradient accumulation.
EMA weights, if you keep an exponential moving average for evaluation.

Verify the resume before you need it

Test recovery on purpose, in a small run, before the expensive one:

Train 200 steps with a fixed seed and record the loss at each step.
Kill the job at step 100 with SIGKILL (not a clean shutdown — the failure you care about is not polite).
Restart with --resume auto and compare steps 101-200 against the uninterrupted baseline.

With full state restored and deterministic settings on, losses should match to within floating-point noise; a divergence at exactly the resume point tells you something in the list above is missing. Then repeat the test at a different world size to prove resharding works. Ten minutes of this saves days later.

Checklist

Compute your interval from measured C and M rather than picking a round number.
Use dcp.save/dcp.load with a Stateful app-state object; never gather the full model to rank 0.
Switch to dcp.async_save and block on the previous future before starting the next.
Write to a temp path, barrier, rename, drop _COMPLETE, barrier; resume only from complete checkpoints.
Save dataloader position, RNG, scheduler, scaler and EMA — not just weights and optimizer.
Launch under torchrun with an --nnodes range and --max-restarts, with --resume auto on by default.
Rehearse a SIGKILL recovery, and a recovery at a smaller world size, before the long run starts.

Fault tolerance is not exotic infrastructure; it is roughly a hundred lines of careful code that converts a class of expensive incidents into a fifteen-minute gap in a loss curve. If your team is about to commit six figures of compute to a single run, it is the cheapest insurance available.

If you would like a second pair of eyes on a training stack before a large run — checkpointing, sharding strategy, throughput or restart behaviour — get in touch; reviewing and hardening distributed PyTorch training is a large part of what our consultants do.

Fault-Tolerant Training Runs: Async Checkpointing with torch.distributed.checkpoint