HomeJournalThis post

Distributed Checkpointing That Recovers

Make sharded training state atomic, topology-independent, corruption-safe, and measurable from failure through the first resumed step.

JP
JP Casabianca
UI/UX designer and full-stack engineer · Bogotá

Distributed checkpointing is reliable only when a training job can change topology, fail mid-save, and resume without mixing incompatible state. This guide designs a manifest-first protocol for sharded tensors, training moments, randomness, and publication.

The intended reader operates multi-rank model training. You will leave with a two-phase save sequence, resharding matrix, corruption fixture, and executable restore gate that treats recovery time as a production objective.

The operating vocabulary connects sharded checkpoint, elastic training recovery, optimizer state, and topology-independent restore within one durability contract.

distributed checkpointing: rank shards converging on an atomic checkpoint manifest An original editorial diagram connects Ranks, Shard writes, Manifest, Resharded restore as one inspectable method. manifest2 ranks
  1. Ranks
  2. Shard writes
  3. Manifest
  4. Resharded restore
Figure 1: Ranks write immutable shards first; only a complete manifest makes the checkpoint visible to a differently shaped restore job.

Distributed checkpointing defines one logical step

The first useful move is to make a globally agreed training-step identity and state inventory visible before choosing an implementation. distributed checkpointing becomes tractable when the inputs and the acceptance line can be inspected together. That framing also prevents a polished demo from answering a different question. The PyTorch Distributed Checkpoint documentation defines planner-based sharded save and load operations intended to support distributed state and topology changes.

Work through four concrete moves:

  • Barrier at an allowed save boundary
  • Record model and optimizer revisions
  • Capture scheduler and scaler state
  • Capture RNG streams and data cursor

I would begin with the smallest representative specimen, then add one difficult edge case and one intentionally broken control. The specimen makes the mechanism legible; the edge case tells us where it bends. The broken control proves the test can reject something. The local check is comparing step identity across every rank.

The failure to watch is letting ranks label shards from different steps. It matters because an attractive average can conceal the exact cohort, state, or frame that makes the method unsafe.

Use no write begins until logical ownership agrees as the decision rule. Preserve the inputs, output, and rejected control together so another person can rerun the claim. That compact receipt is more useful than a universal best practice.

Describe every shard in a manifest

Treat tensor name, global shape, dtype, placement, byte range, checksum, and writer as the working material, not as setup that disappears behind a result. In distributed checkpointing, the shape of that material determines which comparison is honest. Write it down before tuning anything. The Megatron Core distributed checkpointing guide documents sharded state dictionaries, strategies, and fully parallel save and load behavior in large-model training.

Work through four concrete moves:

  • Use immutable object names
  • List expected shard count
  • Store schema and framework versions
  • Include non-tensor training state

Run the sequence once by hand before automating it. A hand-worked example exposes units, ownership, and ordering mistakes that disappear inside a dashboard. Automation should preserve that explanation, not replace it. The local check is reconstructing the full inventory without workers.

The failure to watch is treating a directory listing as a format. It matters because an attractive average can conceal the exact cohort, state, or frame that makes the method unsafe.

Use a self-contained versioned manifest as the decision rule. Preserve the inputs, output, and rejected control together so another person can rerun the claim. That compact receipt is more useful than a universal best practice.

Publish atomically after all writes

A reproducible study starts by isolating a private write phase followed by one durable commit marker. This gives distributed checkpointing a stable object to measure and a clear place for creative judgment. Without that anchor, every later improvement can be explained away by a changed input.

Work through four concrete moves:

  • Write shards to temporary identities
  • Verify size and checksum
  • Collect rank acknowledgements
  • Commit the manifest in one operation

Keep the raw observation beside the transformed result. This makes aesthetic choices discussable and engineering claims falsifiable. It also gives the next iteration a known starting point instead of a screenshot with no provenance. The local check is killing one writer before publication.

The failure to watch is exposing partial saves to restore discovery. It matters because an attractive average can conceal the exact cohort, state, or frame that makes the method unsafe.

Use only committed manifests are listable as the decision rule. Preserve the inputs, output, and rejected control together so another person can rerun the claim. That compact receipt is more useful than a universal best practice.

SignalChoiceEvidence
8 → 8 ranksControl92 s restore
8 → 4 ranksRequired118 s · exact
Missing shardRejectNo state exposed
Figure 2: A checkpoint is accepted only after same-shape, resharded, and corruption paths produce explicit outcomes.

Reproduce a topology-changing restore

Here the design problem is a planner that maps global tensors from eight writers onto four readers. It is both technical and editorial: the system needs a reliable constraint, and the reader needs to see why that constraint matters. Good distributed checkpointing keeps those two views aligned.

Work through four concrete moves:

  • Load and validate the manifest
  • Calculate new shard ownership
  • Read only required byte ranges
  • Compare reconstructed state hashes

Use a narrow worksheet with one row per decision. Name who owns the row, what can change it, and what evidence closes it. This turns critique into a concrete comparison instead of a preference contest. The local check is running the included compatibility fixture.

The failure to watch is assuming rank numbers are permanent tensor owners. It matters because an attractive average can conceal the exact cohort, state, or frame that makes the method unsafe.

Use global coordinates determine restore placement as the decision rule. Preserve the inputs, output, and rejected control together so another person can rerun the claim. That compact receipt is more useful than a universal best practice.

Runnable artifact. Save this inspectable specimen as distributed-checkpointing.test.mjs and run node --test distributed-checkpointing.test.mjs. Expected result: PASS: incomplete checkpoints stay private.

import assert from "node:assert/strict";
import test from "node:test";
const publish=({expected,written,checksums})=>written===expected&&checksums.every(Boolean);
test("manifest commits complete shards only",()=>{assert.equal(publish({expected:8,written:8,checksums:Array(8).fill(true)}),true);assert.equal(publish({expected:8,written:7,checksums:Array(7).fill(true)}),false);console.log("PASS: incomplete checkpoints stay private")});

Preserve optimizer and randomness

The first useful move is to make moment tensors, parameter groups, loss scaling, scheduler step, RNG, and sampler cursor visible before choosing an implementation. distributed checkpointing becomes tractable when the inputs and the acceptance line can be inspected together. That framing also prevents a polished demo from answering a different question.

Work through four concrete moves:

  • Key optimizer slots by parameter identity
  • Record every RNG generator
  • Persist consumed-data position
  • Resume one deterministic step

I would begin with the smallest representative specimen, then add one difficult edge case and one intentionally broken control. The specimen makes the mechanism legible; the edge case tells us where it bends. The broken control proves the test can reject something. The local check is matching the next loss and batch identifiers.

The failure to watch is checking only model weights. It matters because an attractive average can conceal the exact cohort, state, or frame that makes the method unsafe.

Use a one-step continuation receipt within tolerance as the decision rule. Preserve the inputs, output, and rejected control together so another person can rerun the claim. That compact receipt is more useful than a universal best practice.

Bound storage and network pressure

Treat save interval, dirty bytes, staging memory, bandwidth, retention, and deletion as the working material, not as setup that disappears behind a result. In distributed checkpointing, the shape of that material determines which comparison is honest. Write it down before tuning anything.

Work through four concrete moves:

  • Measure checkpoint bytes
  • Throttle background writes
  • Avoid synchronized storage bursts
  • Retain independent recovery points

Run the sequence once by hand before automating it. A hand-worked example exposes units, ownership, and ordering mistakes that disappear inside a dashboard. Automation should preserve that explanation, not replace it. The local check is training under a real save cadence.

The failure to watch is benchmarking saves while the job is idle. It matters because an attractive average can conceal the exact cohort, state, or frame that makes the method unsafe.

Use a tail-step slowdown and recovery-point objective as the decision rule. Preserve the inputs, output, and rejected control together so another person can rerun the claim. That compact receipt is more useful than a universal best practice.

  1. FreezeFreeze

    Capture one logical training step and ownership map.

  2. WriteWrite

    Persist immutable rank shards and checksums.

  3. PublishPublish

    Commit a complete manifest atomically.

  4. RestoreRestore

    Validate, plan resharding, and resume once.

Figure 3: Publication is the commit point separating private partial shards from one discoverable recovery unit.

Test corruption and interruption

A reproducible study starts by isolating missing, truncated, stale, duplicated, and checksum-failing shards plus coordinator loss. This gives distributed checkpointing a stable object to measure and a clear place for creative judgment. Without that anchor, every later improvement can be explained away by a changed input.

Work through four concrete moves:

  • Inject one fault per run
  • Reject before mutating live state
  • Preserve the prior checkpoint
  • Clean private partial objects safely

Keep the raw observation beside the transformed result. This makes aesthetic choices discussable and engineering claims falsifiable. It also gives the next iteration a known starting point instead of a screenshot with no provenance. The local check is a table of terminal outcomes.

The failure to watch is retrying a corrupt shard until it appears valid. It matters because an attractive average can conceal the exact cohort, state, or frame that makes the method unsafe.

Use fail closed with the last committed point intact as the decision rule. Preserve the inputs, output, and rejected control together so another person can rerun the claim. That compact receipt is more useful than a universal best practice.

Measure recovery time to useful training

Here the design problem is detection, provisioning, metadata load, data transfer, resharding, validation, and first successful step. It is both technical and editorial: the system needs a reliable constraint, and the reader needs to see why that constraint matters. Good distributed checkpointing keeps those two views aligned.

Work through four concrete moves:

  • Start the clock at failure
  • Include replacement capacity delay
  • Record phase timings
  • Verify exactly one resumed step

Use a narrow worksheet with one row per decision. Name who owns the row, what can change it, and what evidence closes it. This turns critique into a concrete comparison instead of a preference contest. The local check is a scheduled full-cluster recovery drill.

The failure to watch is reporting raw storage throughput as recovery. It matters because an attractive average can conceal the exact cohort, state, or frame that makes the method unsafe.

Use an end-to-end recovery-time objective as the decision rule. Preserve the inputs, output, and rejected control together so another person can rerun the claim. That compact receipt is more useful than a universal best practice.

Make the method yours

Make the manifest the checkpoint's atomic commit. Immutable shards, checksums, global tensor coordinates, and complete non-tensor state let a restore job reject partial work before touching live memory.

Prove the format by restoring onto a different topology and continuing one deterministic step. Recovery time, not save throughput alone, tells you whether the design protects expensive training.

Continue through four related field notes: durable AI agent execution, data backfill evidence, Postgres worker queues, advisory-lock leaders. They extend the same craft without changing this article's single search intent.