Unit 8.7 · about 5 h
227 / 324Complete checkpoints, atomic writes and reproducibility boundaries
Resume bugs waste entire training runs: the loss looks fine while the data order, the schedule, or the augmentation stream quietly diverges. The checkpoint contract is what makes a restart a continuation instead of a new experiment.
- implement
You can now
- Save and restore every state a resume needs: model, optimizer, scheduler, scaler, epoch, RNG and data position.
- Write checkpoints atomically so a crash never leaves a half-written file on the load path.
- State the exact reproducibility boundary: what matches bitwise, what matches statistically, and what never matches.
Covers
- Complete checkpoint and resume state
- Atomic writes and reproducibility boundaries
Study, in this order
3 steps
- General checkpoints for resumePyTorch - reliability and performance recipes
What is a
state_dict; Saving and Loading a General Checkpoint:saveandloadpatterns, optimizer state, 2 to 3 times model size,.tarconvention,evalversustrainmode; thebest_model_statereference pitfallDepth · Skip Saving an Exported Program and the
DataParallelsection unless your model uses them. - Reproducibility: seeds, determinism and data loadingPyTorch - checkpointing, numerical accuracy and CUDA semantics
Controlling sources of randomness including
manual_seedacross devices; Avoiding nondeterministic algorithms with the cross-release and CPU-versus-GPU limits;DataLoaderworker seeding withworker_init_fnandgenerator; CUDA convolution benchmarking versus determinism - Atomic renames for checkpoint writesPython os - `os.replace`
os.replace: a successful rename is atomic, so write the checkpoint to a temp file in the same directory,flushandfsync, then replace the final name
Focus on
- A resume checkpoint is a dict of state, not a model file: epoch, model, optimizer, scheduler, scaler, RNG states and the sampler position all belong in it.
state_dictreturns references to live tensors. Training on after building the dict mutates the saved copy too; serialize to disk ordeepcopyat save time.- Per-worker RNG needs explicit construction. A
Generatorbuilt inDatasetinit is fork-copied identically into every worker and repeats its stream. - Exact resume is realistic only on one machine, build and device kind, with deterministic algorithms. The reproducibility notes promise nothing across releases, platforms, or CPU versus GPU, so declare a tolerance there.
Reproduce closed-book
torch.equal held on CPU with PyTorch 2.14). Rerun without restoring the RNG and show the weights differ.Traps
- Reloading from the in-memory dict after further training. The dict aliases live optimizer buffers, so the resume silently starts from the later state.
- Seeding only
torch.manual_seedwith severalDataLoaderworkers. Global RNGs get per-worker seeds, but a copiedGeneratorobject repeats exactly; build it per worker fromtorch.initial_seed. - Writing the checkpoint straight to the final name. A mid-write crash leaves a corrupt file on the load path; temp file plus
os.replaceavoids it.
Practice
TrenTorch: Resume-from-checkpoint
Implement save_checkpoint, train_n_steps, resume_from_full_checkpoint and resume_from_weights_only around an Adam step
Restoring \(m\), \(v\) and \(t\) reproduces the uninterrupted weights; a weights-only resume restarts Adam's bias correction and diverges. The section mini-project then adds resume equivalence and interruption handling with declared tolerances.
Self-test
-
Name the checkpoint contents for exact resume on one GPU of
SGDwith momentum, aStepLRscheduler,fp16AMP with aGradScalerand a shuffled loader.Model, optimizer with its momentum buffers, scheduler, scaler, epoch and step, CPU and CUDA RNG states, and the sampler position or loader generator state. -
A saved dict is kept in memory, training continues, then a fresh optimizer loads that dict. Why are its buffers wrong?
state_dictreturns references, so continued training mutated the same buffers the dict aliases. Serialize to disk ordeepcopyat save time. -
Two
DataLoaderworkers emit identical augmentation streams. State the bug and the fix.AGeneratorbuilt inDatasetinit was fork-copied identically into both workers. Reseed it per worker inworker_init_fnfromtorch.initial_seed(), or create it lazily inside the worker. -
The job dies mid-write and the checkpoint is corrupt. What write pattern survives this?
Write to a temp file in the same directory,flushandfsync, thenos.replaceonto the final name. Readers only ever see complete files.
Completion criterion
Stop when
Your disk-roundtrip resume matches the continuous run (exactly on CPU, or within a declared tolerance elsewhere), and a changed-seed or unrestored-RNG resume is caught by your own equivalence check.
Mark passed, then open 8.8 Only mark this when you can do it without notes.Leads to
- 3.8 Diagnosing failed runs: divergence, stale state, slices and seeds Use resume equivalence as a control when a restarted run diverges.
Comes after
- 2.8 Checkpoints, configuration and reproducibilityHarden resume with atomic writes and full state audits.
Resources for this unit
- PyTorch - reliability and performance recipes Use the Profiler, Automatic Mixed Precision, model loading and Benchmark recipes, then the Performance Tuning Guide. Apply each to one compact model from an earlier section and keep a CPU fallback when CUDA is unavailable.
- PyTorch - checkpointing, numerical accuracy and CUDA semantics Use the official user-guide entries for torch.utils.checkpoint, numerical accuracy, CUDA semantics, reproducibility and serialization. Predict memory categories before measuring and define resume equivalence before running.
- Python os - `os.replace` Read
os.replacefor the atomic-rename step of checkpoint writes: save to a temp file in the same directory,flushandfsync, then replace the final name.
Also in this section
- PyTorch - compiler troubleshooting Read graph breaks, guards, recompilations and logging. Compile one stable fragment, intentionally trigger a graph break or recompilation, inspect it, and compare eager and compiled outputs and gradients.
- Deep Learning Training Debugger Use the failure-injection, measurement and evidence rules as the deeper standard. The full Tier 2 project is a later route and is not added to this stage milestone.
- TensorTonic - mixed-precision memory and transformer FLOPs Solve only Mixed-Precision Training Memory Accountant and Transformer Training FLOP Estimator. Reconcile both with the compact model's profiler and document every included and excluded term.
- Learn PyTorch - profiler walkthrough Use notebook 10 only if the official profiler route remains unclear. Transfer the relevant annotations into the compact model instead of keeping the tutorial code.
- TrenTorch - precision, checkpointing and resume Solve Loss Scaling, FP16/BF16 Representable Range, Gradient Checkpointing and Resume-from-checkpoint in NumPy, then reproduce each behavior with the PyTorch API in the compact model.
- Hands-On Machine Learning with Scikit-Learn and PyTorch Read Appendix B, Reduced Precision Models and Mixed-Precision Training, for why
float16training underflows and whatfp32primary weights and loss scaling fix.