Sharded Checkpoint and Atomic Resume

A vanilla checkpoint gathers all parameters and optimiser state onto rank 0 and writes a single file. For a 70B model that is 1.1 TB of state through one rank's network port, and the other ranks sit idle until the gather finishes. Write bandwidth is the slowest single GPU's network link, not the aggregate. On a real cluster the gather-then-write step can take longer than the previous training hour, so the job ships less than one checkpoint per training day.

A sharded checkpoint flips the pattern: each rank writes its own shard to its own file in parallel, and a manifest records which rank owned which shard so resume can put each one back where it came from. Aggregate write bandwidth scales with the cluster. A 1 TB checkpoint that took 4 hours through one rank takes 4 minutes through 64 ranks. The manifest also gives you a contract for incompatible resumes. A world-size change is detectable, a partial write is detectable, and the load path can fail loudly instead of quietly using stale data.

A sharded checkpoint is a checkpoint in which each rank writes its own slice of the training state to its own file. A manifest is a small JSON file that records which shard belongs to which rank, plus the world size and the per-shard hashes needed to validate the shards on load. An atomic write puts a file at its final name in one step, so a reader sees either the complete old file or the complete new one.

Manifest schema

Abridged to two of the four shards:

{
  "world_size": 4,
  "step": 1234,
  "wall_clock_seconds": 4521,
  "shards": [
    {"rank": 0, "path": "rank0.bin", "sha256": "...", "param_shard_offset": 0, "param_shard_numel": 65536},
    {"rank": 1, "path": "rank1.bin", "sha256": "...", "param_shard_offset": 65536, "param_shard_numel": 65536}
  ],
  "schema_version": 1
}

Three fields are load-bearing. world_size makes a resume on a different size fail loudly instead of silently corrupting state. sha256 per shard catches partial or corrupted writes. param_shard_offset and param_shard_numel record where each rank's slice sits in the flat parameter tensor (the offset is a running sum of the earlier shards' sizes). The demo's load_sharded does not read them and simply indexes by rank, but a cross-world-size loader (exercise 4) needs them to concatenate the slices and re-shard.

Atomic write

The standard pattern is to write each shard to <name>.tmp, write the manifest to manifest.json.tmp, fsync each, then os.replace each shard into place and the manifest last. POSIX rename within the same filesystem is atomic, so each file is either fully new or fully old. That guarantee is per file, not per checkpoint: save_sharded does one os.replace per shard and then one for the manifest, all into the same directory. A crash before the first os.replace leaves only .tmp files and the previous checkpoint untouched. A crash partway through the replaces leaves a mix, and the load checks catch it: a missing manifest if the directory is new, or a sha256 mismatch on the first replaced shard if the directory held an older checkpoint. Writing each step into its own directory (step_0100) keeps a fresh save from ever overwriting the live one. Without atomic write, a crash can leave a partial shard next to a manifest that points at it, and the load corrupts the optimiser state on resume.

Three failure modes the schema must defend against

Each defence rejects the bad load early. The alternative is silent corruption that surfaces 100 steps later when loss goes to NaN.

The loader also rejects a schema_version it does not recognise, and any shard path that is absolute, contains a separator, or resolves outside the checkpoint directory, so a tampered manifest cannot point the loader at another file. Shards are read with torch.load(..., weights_only=False), which unpickles arbitrary objects. The sha256 is stored in the manifest next to the data it covers, so it detects corruption, not a malicious writer. Only load checkpoints you wrote.

Why per-rank files, not one big file

Concurrent writes to one file via O_APPEND work on POSIX for byte-aligned writes, but in practice the offsets within one shard span MB-sized regions and locking dominates. Per-rank files have no contention and benefit from striping when the underlying filesystem is parallel (Lustre, GPFS). Production stacks such as DeepSpeed, FSDP, and NeMo use per-rank files for that reason.

code/main.py implements:

• ShardEntry and ShardManifest dataclasses with the schema above plus to_json/from_json.
• save_sharded(per_rank_state, dest_dir, step, wall_clock_seconds=0.0) that writes each rank's binary state to its own .tmp file, then renames the shards and the manifest into place, manifest last.
• load_sharded(src_dir, expected_world_size) that validates the manifest, verifies each shard's sha256, and returns the manifest and the per-rank state dicts.
• rotate_checkpoints(parent_dir, keep_last=5) that deletes the oldest step_* directories.
• A round-trip test: build per-rank state, save, load, and torch.equal each rank's param_shard. The m_shard and v_shard optimiser tensors are saved and loaded with it but not compared.

Run it:

python3 code/main.py

The demo runs four steps in order: write 4 shards plus a manifest and reload them with the parameter shards checked equal, reload with the wrong world size (8) and watch it get rejected, append bytes to rank0.bin and watch the sha256 check reject it, then write 8 more checkpoints (step_0000 to step_0007) next to step_0100 and rotate down to the newest 5. That makes nine directories, so four are deleted: step_0100, step_0000, step_0001 and step_0002.

Three patterns harden the checkpoint enough to ship.

Async write. Production stacks issue the checkpoint write on a separate thread or process so training continues. The barrier sits at the next checkpoint: the next save must wait until the previous one is complete. DeepSpeed's async_io flag handles this. code/main.py keeps the write synchronous so the steps are visible.

Local fast disk first, then async upload. Write to local NVMe, which is fast, then async-upload to S3 or GCS. The two tiers keep the in-cluster checkpoint quick to resume from while a durable copy ships off-cluster for archive. The manifest carries the local path, and an upload manifest carries the remote path.

Rotation matters. Production runs keep the last K checkpoints (typically 3-5) and rotate the oldest out. Without rotation the disk fills mid-run and the next checkpoint fails. With rotation, old checkpoints get deleted and the budget frees up. The demo writes all eight checkpoints and then rotates, so it briefly holds nine directories against K=5. Rotating after every save would peak at K+1, and exercise 2 moves the rotation before the save.

Where the pattern shows up:

• DeepSpeed checkpointing. deepspeed.save_checkpoint(tag=step) writes per-rank files and a latest file pointing at the active tag.
• PyTorch FSDP checkpointing. torch.distributed.checkpoint saves sharded state with a Planner that decides the per-rank layout.
• NeMo. Wraps DeepSpeed and FSDP behind a uniform save_to_checkpoint API that adds metadata.

Lesson 81 saves a sharded checkpoint of the end-to-end DDP+ZeRO run and reloads it on the same world size to show that the resume contract holds. This lesson's main.py stops at the reload: no optimiser step runs after the load, so the optimiser shards are restored but never exercised.

1. Add async write: kick off the save in a thread and let training continue. Block the next save until the previous one completes.
2. Extend the rotation in main.py. rotate_checkpoints(parent_dir, keep_last=5) already exists and main() calls it after writing eight checkpoints, ordering them by directory mtime. Make it run before the new save with keep_last set to K-1, so disk use peaks at K checkpoints mid-write instead of K+1, and order by the step in each manifest instead of mtime.
3. Add a CRC-only fast verification path for the inner-loop reload (rotation rolls a checkpoint into being the new active one without full sha256).
4. Add a cross-world-size load: shard rebalance from N=4 to N=8 by reading the manifest, concatenating, and re-sharding.
5. Add an upload to a fake S3 (a second directory) and write the upload manifest. Defend the two-tier storage policy.

• DeepSpeed checkpointing
• PyTorch torch.distributed.checkpoint
• POSIX rename atomicity
• Phase 19 Lesson 78 - the ZeRO state this checkpoint is shaped to save
• Phase 19 Lesson 81 - the end-to-end demo round-trips the saved state