Checkpoint Save and Resume

A checkpoint is a single file that holds everything needed to continue a run: model parameters, optimizer state, scheduler state, the loss history for plots, the current step and epoch and batch-in-epoch counters, and the RNG state for every source of randomness.

Here is what goes wrong without one. You set a training job for 18 hours. The wallclock cap is 4 hours. The cluster reboots at hour 11 because someone above your pay grade approved a kernel upgrade. With no checkpoint you start over. With weights but no resume logic you still lose the optimizer state that took the first 11 hours to learn, so the AdamW moments are gone and the next step lurches in a direction the training trajectory had already moved past.

RNG state matters too. Leave it out and the resumed loss curve is a different curve: same model, same data, different shuffle, different dropout mask, different number on the dashboard.

Atomic save is the other half of the contract. Write straight into the final filename and a crash mid-write leaves a corrupt file, so the resume reads garbage. Write into a temporary file in the same directory and rename it, and a crash mid-write leaves the previous good file untouched. The rename is atomic on POSIX file systems.

A checkpoint is the full training state serialized to disk, split into five buckets that together let a fresh process pick the run back up.

The five state buckets

Atomic save

An atomic save writes the payload to a temporary file and renames it over the target, so the target name points at a complete file.

Two rules apply. The temporary file lives in the same directory as the target, so the rename stays within one file system (cross-device renames are not atomic). And the temporary name is unique per attempt, so two writers do not stomp on each other.

Sharded checkpoints

A sharded checkpoint splits the parameter state into several shard files and writes a small index that ties them together. You want one once the model is large enough that a single-file payload is too big to load fast, too big to inspect, and too painful when a network share hiccups mid-read.

The index records the shard count, the sha256 of each shard, and the sha256 of the meta file. The loader checks the meta hash first, then each shard's hash as it reads that shard, and fails loudly on any mismatch. In main.py all shards are written to one directory, but because each shard is its own file you could place them on different disks or hosts. The meta file is small and reads first.

Resume continues mid epoch

A mid-epoch resume picks up at the exact batch inside the current epoch where the run stopped, using (epoch, batch_in_epoch). Snapping to the start of the next epoch instead wastes anywhere from minutes to a day. In main.py each epoch's batches come from a torch.Generator seeded with 12345 + epoch, so after load the training loop replays the batch_in_epoch draws already consumed in the current epoch and continues from there. The global Python, NumPy, and PyTorch RNG state is restored too, which matters for dropout and in-model sampling. The demo model has no dropout, so its comparison does not exercise that part. The demo's assertion is that max_loss_diff_after_resume, the largest absolute difference between the uninterrupted and resumed losses after the interruption step, stays under 1e-4.

code/main.py holds four primitives and a demo driver.

Step 1: capture and restore RNG state

capture_rng_state returns a dict with Python's random.getstate, NumPy's np.random.get_state, and PyTorch CPU and CUDA RNG bytes. restore_rng_state reverses it. The PyTorch states are stored as plain lists (via .tolist()) and rebuilt into uint8 tensors on restore, the byte-buffer form PyTorch's RNG consumes.

Step 2: atomic save

atomic_save writes the payload to a temp file in the target directory, then os.replace swaps it into the final name. atomic_write_json does the same for the sharded index.

Step 3: full checkpoint round trip

save_checkpoint packages the model, optimizer, scheduler, train state, and RNG into one dict. load_checkpoint reverses it and returns a TrainState. The schema field is the upgrade hook. When the format changes, you bump the version string and the loader dispatches on it.

Step 4: sharded variant

save_sharded_checkpoint round-robins the sorted parameter keys across N shards, writes each shard with its own atomic save, writes a meta file with optimizer and scheduler and train state, and writes the JSON index with the shard sha256s and the meta file's sha256. load_sharded_checkpoint checks the meta hash from index.json up front, then for each shard verifies its hash (taken from the shard list in the meta file) and merges that shard's tensors before moving to the next.

Step 5: resume demo

run_resume_demo trains a small model for total_steps, saves a checkpoint at interrupt_at, then continues to total_steps to produce the uninterrupted baseline. Next, in the same process, it reseeds, builds a fresh model, optimizer, and scheduler, loads the checkpoint into them, and runs the remaining steps. The function returns max_loss_diff_after_resume, the max absolute difference between the two loss trajectories after the interruption point. The match comes from the epoch-seeded batch generator plus the restored optimizer and scheduler state, so the difference is zero or floating-point noise.

Run it:

python3 code/main.py

Both the single-file and sharded demos assert a max diff under 1e-4, and the summary lands in outputs/resume-demo.json.

Production training stacks ship checkpointing as part of the trainer, and the shape matches what you built: model, optimizer, scheduler, counters, and RNG, written atomically (production trainers usually also name files by step so the latest is easy to find; main.py writes a fixed ckpt.pt). Sharded layouts power large model loading with parallel reads, and the index.json is what makes that work.

Three patterns to enforce:

• Schema is a string in the payload. Migrations branch on it. Without it you cannot evolve the format without breaking old runs.
• Sha256 every shard. A silently truncated download is a nasty bug, and the loader either fails fast or fails late.
• Keep checkpoint cadence honest. Save every N steps and every wallclock-minute, whichever is shorter. Otherwise the long step that crashes wastes a full window of work. main.py saves at one chosen step only; the wallclock trigger is Exercise 3.

outputs/skill-checkpoint-save-resume.md is the recipe for a new training script: payload shape, atomic write, RNG capture, sharded index. Drop the skill into a repo, wire save_checkpoint at the periodic save site and load_checkpoint at startup, and the run survives kills.

1. Replace round-robin sharding with sharding by parameter group (layers ending in .weight vs .bias). When is each layout preferable?
2. Extend the save loop to keep the last K checkpoints and prune older ones. What is the right K when the disk is small?
3. Add a --ckpt-every-seconds flag that triggers a save on a wallclock interval, not just step count.
4. Add a checksum verification path that runs at startup, scans every checkpoint in the directory, and reports which ones are corrupt.
5. Implement a migrate_v1_to_v2 function that adds a new field to the payload and bumps the schema string. Make load tolerate both versions.

• POSIX rename semantics for the atomicity claim that os.replace relies on.
• PyTorch documentation on torch.save and torch.load, including map_location for cross-device restores.
• Phase 19 lesson 46 covers the gradient accumulation that the checkpoint payload here survives across.
• Phase 19 lesson 48 covers the distributed wrappers whose state dict format this scheme accommodates.
• The Linux kernel fsync documentation for the durability guarantee behind atomic rename.