explanationView source

Distributed checkpoints and resume#

A distributed checkpoint has to answer a harder question than a single-device one: is this one consistent state of one learner, or a mixture of owners that each saved at a slightly different moment? In Alpha a checkpoint is published only when every owner's piece belongs to the same capture, and the published root names every piece by digest. A resumed run either restores exactly that state or refuses before any owner trains.

This page covers how state is saved, published, restored and repartitioned across owners, and what is refused.

What a checkpoint holds#

Checkpoints are indexed by logical identity, never by rank, path or address. Checkpoint.RegionManifest defines the identity of a logical region as two things: the digest of the complete logical schema, and the region's declaration ordinal. Names and formats describe the region, but no rank, path, stream offset or device address is part of its identity.

PartHolds
ShardOne logical region, or an interval of one, from one owner's capture
CaptureThe membership epoch, the committed update number and the input cursor
ManifestThe content digest of every shard, and exactly one capture

Coverage of the logical state is checked on its own terms, not by adding up byte counts. The manifest planner refuses:

  • a wrong format;
  • incomplete or overlapping regions;
  • captures from different moments.

Checkpoint.NativeRegions is the native publisher and restore planner. Restore validates every file before copying any state into its destination, and admits:

  • joins, and splits across serialized words;
  • changed region boundaries;
  • swapped physical offsets.

Paths are runtime bindings, so a relocated set of immutable files still restores.

Publishing durably#

One shared owner, nativeFileCommitStages in Runtime.NativeFilePlan, performs the last four steps of every durable write:

  1. sync the data;
  2. close the file;
  3. rename it into place (with NOREPLACE, so a published generation is immutable);
  4. sync the directory.

The first two stages, the header and the verified payload, belong to the format.

A failed final rename keeps the previous generation, and no participant reports success. The distributed-checkpoint gate kills participants with SIGKILL around the rename to check this.

On a star#

Owner 0 publishes:

  1. Every other owner sends its capture.
  2. Owner 0 checks that all captures agree.
  3. Owner 0 prepares the shards, publishes the manifest, and sends a durable-completion acknowledgment.
  4. Every participant waits for that boundary before reporting the update as committed.

For replicated learners, owner 0 first compares every owner's replicated state with its own, so diverged replicas publish nothing. The owner here is Checkpoint.NativePeerPublication.

On a hierarchy: the manifest tree#

On a bounded-fan-out hierarchy, Checkpoint.NativeTreePublication publishes without any participant holding another participant's tensor bytes. Each participant:

  1. writes its own shards durably to the admitted store;
  2. writes a group manifest: the capture prefix, a record (descriptor, digest) for each of its own shards, and one record per child naming the child's role, its subtree size, and the digest of the child's group manifest;
  3. sends its parent only a commitment: the capture prefix plus that one record.

A parent's work and memory are therefore bounded by the fan-out. The root's group manifest is the published root, and it identifies one complete state through a chain of digests, each level checked by the participant above it.

Coverage is structural. Roles are numbered in preorder, so a child's subtree is the range [child, child + size). A parent checks that each commitment names exactly the child and subtree size the hierarchy prescribes, which by induction establishes complete coverage at the root. Three failure cases follow from this:

  • A missing subtree cannot commit: its parent waits, then aborts.
  • A commitment with a different (version, cursor, epoch) prefix is refused.
  • A commitment that claims one role too many is refused.

Completion is released down the tree only after the root is durable, so no participant reports success for a generation that was never published.

Restoring#

On a hierarchy, nativeTreeRestoreManifests reads the previous generation along one path. It reads the root manifest, then each group manifest down to this role's own, checking each against the record its parent holds and every capture prefix against the root's. It then reads this role's own shards.

A restore touches height + 1 manifests and the role's own shards, never a sibling subtree. Replicated words are read from the root's shards.

Continuation restores the learner state together with its counters: the committed update, the input cursor, and the producing epoch. The membership epoch advances exactly once per continuation. A saved generation from a different epoch is rejected, and left unmodified.

A resume refuses before any owner trains if it finds any of these:

  • a reset cursor or a stale epoch;
  • a missing shard;
  • a NaN or infinite value in the saved state, even with its checksums repaired;
  • another learner's checkpoint;
  • a corrupted shard;
  • a checkpoint saved under a different numerical contract, combine tree or learner;
  • on a GPU, the wrong device UUID or product.

Repartitioned restore#

A saved state can resume on a different number of owners, when its checkpoint was saved without binding the placement.

The key is dataParallelLearnerIdentity in Learning.DataParallel. It covers the learner's rule (programs, initial state, rate) and the input sequence in owner order. It deliberately excludes the split of samples among owners, the owner count and the combine order. A checkpoint that binds this identity is a state of the same learner on any placement.

To restore onto another owner count:

  1. dataParallelRepartitioned splits the same input sequence into contiguous shares for the new count: sizes differing by at most one, larger shares first.
  2. Each new owner restores the shared words, and its own coordinates' words, from the per-word shards.

The parameter-partitioning gate exercises this:

  • Momentum saved on four owners resumes on three, and then on five.
  • Adam saved on four resumes on two. The moments and bias-correction powers move to their new owners.
  • Each result equals an independent oracle of that exact sequence of placements, bit for bit.

The host-only region-checkpoint gate covers one-to-two and two-to-one region mappings.

Only checkpoints with partitioned parameters bind the placement-independent identity today. Replicated and optimizer-state-only checkpoints still bind their placement, because their checkpoints name the star's combine order.

Resume equivalence#

Every learner is held to the same test: five uninterrupted updates must end in exactly the state reached by three updates, a checkpoint, and two more updates in fresh processes.

LearnerWhereResult after five updates
Samples-partitioned scalar learnerTwo RTX 3070sParameter 3/32 (0.09375); prediction from the resumed state 3/16
Coordinate learner, host realizationx86-64 hostsWeights [65059/32768, −16861/16384], from an independent rational oracle
Data-parallel learnersx86-64 and emulated AArch64Bit-identical to the oracle, including after resuming with every owner on the other instruction set

On GPUs the qualification driver also corrupts a saved checkpoint and confirms that the fresh process refuses it (checkpoint-chunk).

Gates#

GateChecks
distributed-checkpointPublication, fresh restore, corruption, and SIGKILL around the rename
region-checkpointShard manifests and repartitioned host restore
hierarchical-learningManifest trees at fan-out 1, 2 and 3
parameter-partitioningRestore onto three, five and two owners
numerical-contractsA changed contract rejects a valid checkpoint

Not built yet#

  • Repartitioned restore on a hierarchy, and for replicated or optimizer-state-only checkpoints.
  • A star restore that reads only the shards an owner needs. Today the region restore validates every shard in the manifest.
  • Failure-tolerant publication authority beyond one root publisher, and placement of shards across failure domains.
  • Power-loss durability testing. The store is the POSIX filesystem named by the paths.
  • Qualifying on real hardware: resume equivalence on physical GPUs.
  • Toward 200,000 devices: why the manifest tree is shaped the way it is.