A successful write proves only that a writer returned
A training job can log checkpoint complete even when the result cannot resume training. One rank may have failed to publish its shard, a pointer may reference an incomplete directory, or the file may contain model weights without the optimizer and data position. A checksum can prove that stored bytes match a manifest, but not that the training code can interpret those bytes under the intended topology and software environment.
Treat checkpoint creation and recovery as separate operations with separate service indicators. Creation asks whether all required artifacts were durably published. Recovery asks whether a fresh process can discover a committed checkpoint, load every required state, resume at the intended training position, and produce plausible subsequent steps. The second operation is the evidence that matters during an interruption.
Inventory state by its effect on the next step
Model parameters and buffers are only the obvious layer. Optimizer state can contain momentum or adaptive moments that change the next update. A learning-rate scheduler carries its own position. Gradient scaler state matters in mixed-precision training. Accumulation counters determine whether the next microbatch should produce an optimizer step. Framework-specific parallelism metadata may be necessary to map sharded state back onto ranks.
PyTorch's saving and loading tutorial explicitly says a general checkpoint for resuming training needs more than the model state dictionary and shows model, optimizer, epoch, and loss in one checkpoint. Extend that principle to everything whose value changes future execution. Maintain a required-state registry in code so a new scheduler, scaler, or sampler cannot be added to training without declaring its save, load, and validation behavior.
| State family | Recovery question | Useful validation |
|---|---|---|
| Model and buffers | Did every expected tensor load? | Keys, shapes, dtypes, and sampled checksums |
| Optimizer and scaler | Will the next update use prior history? | State entries, step counters, finite next update |
| Scheduler | Is the learning rate at the intended step? | Per-group learning rate before and after one step |
| RNG state | Does stochastic work continue from the saved stream? | Known post-restore draws per rank |
| Sampler and data loader | Which examples come next? | Next sample IDs and epoch position |
| Manifest | Is this one complete committed checkpoint? | Artifact count, sizes, checksums, and commit marker |
Randomness and sampler progress are training state
Saving only an initial seed is not equivalent to saving generator state after millions of draws. Python, NumPy, CPU tensor operations, and each CUDA process can advance different random streams. PyTorch exposes functions for getting and setting RNG state, while distributed frameworks may maintain additional trackers. Capture the streams actually used and associate rank-specific state with a stable rank or topology identity.
The data path also needs a cursor. Record epoch, sampler permutation or reproducible generator state, consumed sample count, and worker or sharding information needed to reconstruct the next batch. Iterable datasets may require source offsets or durable record identifiers. The validation is concrete: after restoration, the next several sample IDs should match the expected continuation, not silently restart the epoch or skip an unbounded section.
Publish a manifest after every component is durable
Use a manifest that names the checkpoint format version, run and step, creation time, code and configuration identifiers, world size and parallelism layout, model and tokenizer identity, dataset snapshot, and every artifact with byte length and checksum. Write shards to a staging location, verify them, write the manifest, and only then publish a small commit marker or update the active pointer. Recovery should ignore any directory without the committed state.
The manifest should also state which components are intentionally absent. A weights-only export may be valid for inference but must not masquerade as a resumable training checkpoint. PyTorch's general-checkpoint example distinguishes resuming training from loading weights for inference by including optimizer state and training position. Data-loader, RNG, scheduler, scaler, and other state still need explicit treatment when they affect continuation. Names differ by framework, but explicit completeness remains the operable idea.
Worked example: test the next update, not just the load
Consider an illustrative 16-rank job checkpointed after optimizer step 48,000 with four gradient-accumulation microbatches per step. The manifest lists 16 model shards, 16 optimizer shards, one scheduler file, one scaler file, 16 rank-specific RNG files, one sampler file, and one configuration file. A clean recovery worker first verifies all 52 listed artifacts and their checksums, loads them without permissive missing keys, and confirms that the restored scheduler state produces the expected learning rate for the next update. Scheduler counters and step timing differ across implementations, so the manifest must define whether a stored counter names the completed step or the next step.
The drill then reads the expected next sample IDs, runs four microbatches, and performs step 48,001. It checks that loss and gradient norm are finite, the optimizer counter advances once, and the learning rate follows the stored schedule. For a stricter test, compare selected tensor checksums or numerical summaries with an uninterrupted control after the same deterministic step. Exact equality may be unavailable for nondeterministic kernels, so define tolerances in advance rather than choosing them after a mismatch.
Use a tiered gate with measurable failure criteria
Run three tiers. On every save, verify manifest completeness, sizes, and checksums. On a rotating sample, restore in an isolated process and run a smoke step. On a regular cadence and before major upgrades, terminate a disposable training replica and recover it through the real orchestration path on the supported topology. Record discovery time, load time, first-step time, and recovery point objective loss measured in training steps.
A checkpoint family passes only if every required artifact verifies, strict loading reports no unexplained keys, the next data IDs are correct, optimizer and scheduler counters match the manifest, RNG validation passes where determinism is promised, and the smoke step completes with finite diagnostics. Quarantine a family after one failed sampled restore and retain the last known restorable generation. ClusterBid capacity planning should include the temporary GPU headroom required for these drills when recovery assurance is part of the operating requirement.
Test the boundaries where recovery usually breaks
Repeat drills across the changes the recovery promise actually covers: same topology, supported topology changes, a container rebuild, credential rotation, a different node, and restoration from the remote durability tier. Test a missing shard, corrupt checksum, stale active pointer, and incompatible format version. The expected behavior is a loud refusal or an explicitly supported migration, never a silent fallback to partial state or step zero.
Recovery tests have limits. A smoke step does not prove that a week of resumed training will match an uninterrupted run, and deterministic replay may be impractical across hardware or kernel changes. Object-store visibility, bandwidth, and credential failures can also differ during a regional incident. State the supported recovery envelope, keep more than one verified generation, and measure real restore duration. Completion rate belongs on the checkpoint dashboard, but successful clean-room recovery is the stronger readiness signal.