Unit 2.8 · about 5 h
184 / 324Checkpoints, configuration and reproducibility
An experiment you cannot resume is an anecdote. Checkpoints plus config plus seeds turn a run into evidence you can re-enter.
- implement
- measure
You can now
- Save and restore model, optimizer, scheduler, epoch and RNG state, and state the resume tolerance you actually tested.
- Record hyperparameters, data position and generator state so a resumed run continues the same stream, not a similar one.
- Draw the reproducibility boundary: what matches bitwise, what matches statistically, and what never matches across CPU, GPU and workers.
Covers
- Checkpointing model, optimizer, scheduler and scaler state
- Configuration, data position and random generator state
- Reproducibility claims across CPU, GPU and data workers
Study, in this order
3 steps
- Saving and Loading Models: `state_dict`, checkpoints, devicesPyTorch - Autograd mechanics, serialization and reproducibility
What is a
state_dict(model and optimizer dicts); Saving and Loading Model for Inference (Save and Loadstate_dict,model.eval); Saving and Loading a General Checkpoint for Inference and/or Resuming Training (epoch, model and optimizer state, loss); Saving and Loading Model Across Devices (map_location,DataParallelmodule prefix)Depth · Skip Saving an Exported Program, Saving Multiple Models in One File and Warmstarting. The deprecated general-checkpoint recipe redirects to this page.
- Automatic Mixed Precision recipe: saving and resuming the GradScalerPyTorch Automatic Mixed Precision recipe - Saving/Resuming
Saving/Resuming: the scaler
state_dictsaved beside model and optimizer, when to save it, and resuming across runs with and without AMP - Reproducibility: randomness control and nondeterminism limitsPyTorch - Autograd mechanics, serialization and reproducibility
Reproducibility intro (no cross-release, cross-platform or CPU-versus-GPU guarantee); Controlling sources of randomness (
manual_seed, Python and NumPy seeds, cudnn benchmark); Avoiding nondeterministic algorithms (use_deterministic_algorithms, uninitialized memory);DataLoader(worker_init_fn, generator)Depth · Read the section intros and the
DataLoaderpart fully. Skim the per-op catalog (convolution, SDPA backends, RNN); the boundary table is the assignment, not every kernel.
Focus on
- A resume checkpoint holds model, optimizer and scheduler state, the
GradScalerstate when you train with AMP, the epoch and step, and the RNG states. Model weights alone restart the model but not the training. - Config means every number that changes the stream: hyperparameters, data order and position, transform versions, worker count and seeds. Log them with the checkpoint.
- Identical seeds do not promise identical results across releases, platforms, or CPU versus GPU. Promise bitwise match only inside one tested configuration.
- Deterministic flags trade speed for repeatability and can throw on ops without a deterministic path. cudnn benchmark selection is itself a randomness source.
Reproduce closed-book
Traps
- Saving
model.state_dict()and calling it a training checkpoint. Without optimizer, scheduler, epoch and RNG state, resume restarts the optimization, not the run. - Pickling the whole model and loading it with the default
weights_only=True. Since PyTorch 2.6 that default refuses arbitrary pickled objects, andweights_only=Falseruns pickle, which can execute code. Save state_dicts, and load untrusted files only withweights_only=True. - Claiming GPU determinism from a CPU test. Nondeterministic kernels, benchmark selection and worker seeding each break the claim independently.
Practice
Learn PyTorch 01 task 5: save, reload and verify predictions
Exercises, task 5: save state_dict, reload into a new instance, verify predictions match
Extend it: add optimizer state, epoch and torch, Python and NumPy RNG states, then run the interrupted-versus-uninterrupted resume comparison from Reproduce.
Learn PyTorch 05: scripted training with argparse config (tasks 1 to 3)
Exercises, tasks 1 to 3: get_data.py with existing-data check, train.py with argparse for data dirs, lr, batch size, epochs and hidden units, predict.py that loads the saved model
Treat the argparse flags plus seeds as the config record. A run whose flags are not logged does not count as reproduced.
Self-test
-
The resumed weights give the same validation loss as the uninterrupted run at the checkpoint, yet the training curves split within a few updates. Name three state items that could be missing.
Optimizer state such as momentum or Adam moments, scheduler state (so the learning rate is wrong), and the RNG and data-order state (different shuffles and dropout masks). None of them changes the loss of the saved weights. -
Same seed, same code, CPU run matches but GPU run differs in the last decimals. Is anything wrong?
Not necessarily. The note disclaims CPU-versus-GPU bitwise match; different kernels accumulate floating point in different orders. -
use_deterministic_algorithms(True)throws on your model. What does that tell you, and what is the cost of obeying it?One op has no deterministic implementation. Obeying it means replacing the op or staying nondeterministic; deterministic paths are also usually slower. -
A
DataParallelcheckpoint loads with missing-key errors on a single-GPU model. What happened?The keys carry themodule.prefix fromDataParallel. Savemodel.module.state_dict()for device flexibility, or strip the prefix on load.
Completion criterion
Stop when
An interrupted and an uninterrupted run agree within your written tolerance, and your config log plus boundary statement say exactly what would break that agreement.
Mark passed, then open 3.1 Only mark this when you can do it without notes.Leads to
- 8.7 Complete checkpoints, atomic writes and reproducibility boundaries Harden resume with atomic writes and full state audits.
- 3.8 Diagnosing failed runs: divergence, stale state, slices and seeds Use resume equivalence as a debugging control.
Comes after
- 2.6 Datasets, data loaders and samplersCarry loader position and RNG state into checkpoints.
- 2.7 Explicit training and evaluation loops with correct metricsPersist the loop state you can now run.
Resources for this unit
- PyTorch - Autograd mechanics, serialization and reproducibility Read Autograd mechanics, Serialization semantics and Reproducibility. Test leaf behavior, saved tensors, grad modes, state dictionaries and the exact tolerance promised by one interrupted and resumed run.
- PyTorch Automatic Mixed Precision recipe - Saving/Resuming Read the Saving/Resuming section: save
scaler.state_dict()with the model and optimizer state, at the start of an iteration or afterscaler.update(), and how to resume with or without AMP.
Also in this section
- PyTorch - Learn the Basics Complete Tensors, Datasets and DataLoaders, Transforms, Build Model, Autograd, Optimization, and Save and Load Model. Rebuild each essential example in a small src package instead of retaining a tutorial notebook.
- Learn PyTorch - selected framework exercises Complete exercise 00, tasks 1-5 in 01, tasks 1-7 in 02, tasks 5-13 in 03, tasks 1-3 in 05, and exercises 1-2 in 07. Recreate solutions independently and verify behavior against current official docs.
- Deep Learning with PyTorch, 2nd edition - framework foundations Read Sections 3.8-3.9 and 5.5 when tensor storage, strides or autograd need a second explanation, with the Chapter 3 and 5 companion notebooks for runnable cells. Inspect the Chapter 4 and 7 code only for a diagnosed data-representation or loading gap.
- NeetCode ML - PyTorch closing drills After the module gate, solve Pytorch Basics, Layer Normalization, Batch Normalization and RMS Normalization cold. Explain shapes, state and train-evaluation behavior aloud. These do not replace the package and resume test.
- Hands-On Machine Learning with Scikit-Learn and PyTorch Read Chapter 10, Implementing Mini-Batch Gradient Descent Using DataLoaders and Model Evaluation, for the train and evaluate loops and how batch metrics must be aggregated.