Sharded, portable checkpoints
Map rank-local shards to logical tensor keys, publish only complete versions, and verify restore in a fresh process group without one-rank assembly
Before you begin
Follow a linked prerequisite when you cannot explain it from memory. Background skills without links are stated as abilities rather than hidden terminology.
- Chapter 14 fully sharded state ownership
- Ability to list files or objects and verify required parts
- Understanding that a continuation snapshot must correspond to one completed logical step
At completed step 120,000, a 64-rank training job must preserve enough state to continue after every current process is gone. Model weights are only one category. The continuation also needs optimizer tensors, gradient-scaler state, learning-rate scheduler state, random-number-generator state, the completed-step number, and the data position promised by the input pipeline.
Print or record the state dictionary before saving and separate tensor entries from non-tensor continuation state. A tensor entry has a logical key, element datatype, global shape, and live bytes that may be replicated or divided across ranks. Non-tensor entries such as a completed-step number, scheduler counter, or data cursor use a declared schema and serialization but do not have a tensor shape. Gathering every tensor entry on rank 0 would recreate the memory and bandwidth limit that sharding removed. A distributed checkpoint planner instead assigns rank-local tensor writes and stores range metadata separately from the tensor payloads.
Object-store listings or filesystem files provide observable save evidence, but file presence does not prove recoverability. A failed save can leave 63 complete shard files beside one missing shard. The save protocol therefore writes under a unique checkpoint identity, verifies required data and metadata, and publishes a completion marker only after every required part succeeds.
Test the artifact by terminating the original process group, constructing a fresh group, loading destination tensors, and executing at least one optimizer step. Compare the restored logical state, recorded step, next input position, and accepted numerical behavior. A restore at a different rank count requires both tensor resharding support and an explicit policy for batch, randomness, and data ownership.
Define the continuation contract
List the state needed before choosing a file layout. Parameters and optimizer slots are usually necessary. Mixed-precision training can require the loss scaler and higher-precision master weights. Learning-rate and other schedulers need their counters. Random generators can exist in Python, the tensor framework, each device, and data workers. Data loading can require an epoch, sampler state, cursor, or sample identifiers. The completed-step value must agree with all of them.
Exact continuation and statistically valid continuation are different contracts. Exact continuation attempts to produce the same next batches and random choices after restore. Statistical continuation permits a documented change in sample order or random stream while preserving the training objective. State the contract because it determines what must be stored and what the restore test must compare.
Separate logical tensors from storage shards
For a tensor entry, a logical key identifies an item such as a parameter or one optimizer moment. Tensor metadata describes its global shape, element type, and layout, while physical shard records describe which tensor region is stored in each object or file. The range mapping can use tensor offsets and lengths instead of old rank numbers. Non-tensor checkpoint entries follow their own schema and are not described as tensor shards merely because they share the same checkpoint version.
This separation permits load-time resharding. The restore planner compares stored ranges with target ranges and schedules only the intersections. A target shard may read from several source objects, and one source object may serve several target ranks. No process must materialize the complete logical tensor. PyTorch Distributed Checkpoint uses planners and distributed storage readers and writers for this type of multi-rank save and load.
Locate the consistency boundary
A synchronous save can stop training after a completed step, expose a stable state dictionary, write it, and resume. Every rank must identify the same completed logical step and agree that all collectives and optimizer mutations for that step are complete. The pause is simple to reason about but can be long. An asynchronous path first stages a stable snapshot, often from device memory into host memory, and then uploads that snapshot while later training work proceeds. Staging completion and storage completion are distinct events.
Asynchronous saving is correct only if later optimization cannot mutate the staged values. The implementation can synchronize until device copies finish, use separate buffers, or apply copy-on-write rules. The save design must also account for the lifetime and memory cost of those buffers. If a 40 GiB rank-local state needs another 40 GiB host staging area, host capacity and memory bandwidth become part of the checkpoint design.
Publish one completed version
Give each attempt a unique checkpoint identity. Writers store tensor data and metadata under that identity. After all required writes finish, verify expected objects, sizes, and checksums. Only then publish a small manifest or other atomic completeness record that readers use to discover the version. A directory listing alone is not a completeness protocol.
Keep at least one older verified version until the new version passes the required policy. Storage corruption, schema mistakes, or a software regression can make the newest version unreadable even when every write call returned successfully. Retention should also consider correlated storage faults rather than assuming that two names on one failed device are independent copies.
Restore into allocated target state
Many distributed checkpoint APIs load in place. The application constructs the target model and optimizer for the new mesh, obtains their target state objects, and asks the loader to fill them. Target allocation defines the desired shard shapes. Metadata and the load planner determine which stored ranges fill each target.
Restore also needs process-group initialization, device placement, framework and schema compatibility, and enough temporary memory for reads and conversions. If rank identities can change after an elastic restart, stored state must not depend on rank 17 always representing the same host or data partition. Logical ownership and current membership replace that assumption.
Test the continuation path
A checkpoint test should create a version, terminate or isolate the original processes, construct a fresh process group, restore, and execute at least one training step. Compare tensor shapes and selected values, the recorded step, optimizer state presence, data position, and the next loss under the promised continuation contract. A test that only reads metadata does not test resharding, optimizer construction, or device placement.
Run restore drills for topology changes that operations may actually use. Measure restore bandwidth, slowest-rank completion, temporary memory, and time until the first useful post-restore step. Restore-drill measurements feed the recovery-time term in the calendar model; nominal storage throughput alone does not.
Exact replay after a world-size change may require globally indexed samples, counter-based random streams, and algorithms whose reduction and routing order remain defined across the new mesh. If those conditions are absent, classify the result as statistical continuation and test its accepted invariants instead of promising the identical next batch or bitwise-identical optimizer path.
Worked example: Restoring onto a different mesh
A logical 16 GiB optimizer tensor was partitioned evenly across 64 ranks. Training must resume on 32 ranks after capacity changes, using the same contiguous dimension-zero partitioning rule.
- Each old rank stored 16 GiB / 64 = 256 MiB. Each new rank must own 16 GiB / 32 = 512 MiB. For this aligned layout, new rank 0 needs the ranges formerly held by old ranks 0 and 1; new rank 1 needs old ranges 2 and 3; the pattern continues.
- Read the global shape, type, and range metadata. Allocate each 512 MiB target shard as part of the new optimizer state. Plan two 256 MiB range reads for each target rank. More general layouts can split and combine non-aligned ranges, but still need no 16 GiB full-tensor buffer.
- Restore model parameters and every optimizer slot with the same completed-step boundary. Load scheduler, scaler, and random state. Map the stored data cursor into the new set of workers according to the documented exact or statistical continuation contract.
- Verify checksums for stored objects and target range coverage. Confirm that every logical element is written exactly once to the target state. Missing and overlapping ranges are errors, even if all file reads succeeded.
- Execute one step on 32 ranks. Compare the restored pre-step loss and selected tensor values with a reference restore where practical, then verify that optimizer updates and checkpoint counters advance together.
Conclusion: Logical range metadata makes the 64-to-32 change a planned redistribution of tensor regions, while the continuation contract ensures that non-tensor training state changes consistently with those regions.
Common errors
- Saving only model weights during training. Optimizer, RNG, scheduler, and data position may be required for faithful continuation.
- Publishing a completion manifest before every shard is stored and verified. Publish one atomic completeness indicator after all required data is available.
- Assuming a successful write implies successful restore. Regular restore drills are the only credible validation.
- Indexing state only by old rank number. Elastic restarts and topology changes require logical tensor regions and membership-independent data semantics.
- Calling asynchronous saving free. Account for staging memory, copy bandwidth, upload interference, and the time at which the version becomes recoverable.
Reference summary
Distributed checkpointing saves rank-local shards in parallel and records metadata that describes the logical model state. When the storage format, planner, tensor layouts, and application semantics support resharding, restore can map stored tensor ranges to a different process mesh. In-place loading writes into already allocated destination tensors and avoids assembling a full model on one rank.
- Save tensor entries such as model parameters and optimizer moments, plus the non-tensor scaler, scheduler, random-generator, data-loader, and completed-step state required by the continuation contract.
- Describe each tensor with a logical key, datatype, global shape, layout, and stored byte ranges. Describe non-tensor continuation entries with their own schema instead of pretending that every checkpoint item has a tensor shape.
- Stage data under a unique checkpoint identity, verify all required shards and metadata, and publish one completion indicator last.
- Plan writes and reads across ranks so no process must hold or transfer the complete state.
- Define snapshot consistency when asynchronous staging overlaps training and state can mutate after the selected step boundary.
- Retain a known recoverable checkpoint until a newer checkpoint passes an automated restore and continuation test.
Knowledge check
Why should checkpoint state be keyed by logical model identity rather than the rank that happened to own a shard?
Show the answer guide
- Logical keys and global tensor ranges identify model meaning independently of whichever old rank wrote each physical shard.
- After a world-size or topology change, a restore planner can intersect stored ranges with newly allocated target ranges without materializing the full tensor or assuming that rank identities remain stable.
- Non-tensor continuation state also needs logical schema and a common completed-step boundary; physical filenames and old rank numbers are storage details rather than the continuation contract.
Restore a checkpoint onto a new mesh
Save a sharded tensor plus optimizer, scheduler, random, completed-step, and data-position state, terminate the original processes, and restore with a different rank count or simulated target partition.
Evidence to keep: Tensor and non-tensor schema, complete-version manifest, stored-to-target range map, checksums and coverage test, temporary-memory peak, restored values, and one valid post-restore optimizer step.
Checkpoints provide recoverable state after a failure. Recovery also requires progress detection and identification of the failed process, device, network path, storage operation, or invalid software state.
Sources and further reading
The lesson explanations synthesize primary papers and official references. Follow the originals for proofs, exact APIs, architecture details, and version-sensitive behavior.