Unit 3.5 · about 4 h
308 / 324Checkpoints that really resume
A weights-only checkpoint restarts training but not the run: learning rate, moments and data order all jump. Equivalence tests turn resume from a hope into a measured fact.
- implement
- measure
You can now
- Save and load complete training state: model, optimizer, schedule step, scaler, RNG and data order.
- Bind tokenizer, data and configuration identity to each checkpoint so mismatches fail loudly.
- Prove interrupted training matches uninterrupted training to a stated equivalence level.
Covers
- Model, optimizer, scheduler, scaler and RNG checkpoint state
- Tokenizer, data and configuration identity binding
- Interrupted versus uninterrupted resume equivalence
Study, in this order
5 steps
- CS336 Assignment 1 handout: `checkpointing` and the training loopStanford CS336 Assignment 1 - training and decoding contract
Section 5.2, Checkpointing (p. 35): model and optimizer
state_dictplus iteration throughtorch.saveandtorch.load; then Section 5.3, Training loop (p. 36): configurable hyperparameters, memmap datasets, checkpoint paths and step plus wall-clock logging - LLMs from Scratch ch05: saving and loading weightsLLMs from Scratch companion - training and generation
Section 5.4 (
torch.saveonstate_dict,load_state_dictinto a fresh model, and why adaptive optimizers need their state saved too) - PyTorch saving and loading: general checkpoint plus RNG controlPyTorch saving, loading and reproducibility references
Saving and Loading a General Checkpoint for Inference and/or Resuming Training (Save and Load: model plus optimizer state dicts with the epoch or step)
- PyTorch reproducibility notes: RNG and `DataLoader` statePyTorch saving, loading and reproducibility references
Controlling sources of randomness (torch, Python and other libraries) and
DataLoader(generator seeds for worker shuffling), plus Avoiding nondeterministic algorithmsDepth · Read for the checklist of what counts as state; skip CUDA kernel backend options until GPU work needs them.
- Modern LLM Notebook 09-model-config: repository and `config` identityModern LLM Notebook - configuration, loss and generation
Repository File Map (
config,tokenizer_configandgeneration_configroles) and thetokenizer_configsection (special tokens, max length, truncation and padding sides)
Focus on
- The iteration number restores the schedule exactly, the optimizer state restores the moments, and the scaler state restores the loss scale.
- RNG seeds and data-loader order are training state. A fresh shuffle on resume silently starts a different run.
- Equivalence needs a stated level: exact loss match for a few steps on fixed hardware, then tolerance bands beyond that.
- Tokenizer and config versions pin the meaning of every saved tensor. A changed vocabulary or width redefines all IDs and shapes.
Reproduce closed-book
save and load for model, optimizer, iteration, scaler and RNG seeds, then an interrupted-versus-uninterrupted equivalence test over fixed steps.Traps
- Saving weights only. The resumed run restarts the schedule and moments, so loss jumps even though the model is identical.
- Recreating the DataLoader without its generator state. Batches arrive in a new order and the run diverges from the uninterrupted reference.
- Loading weights into an edited config. Shapes may still match while token IDs or widths mean something new, and nothing raises an error.
Practice
CS336 Assignment 1: `checkpointing`
Problem (checkpointing) in Section 5.2; connect save and load to run_save_checkpoint and run_load_checkpoint in tests/adapters.py and pass test_checkpointing in tests/test_serialization.py
The test trains 10 AdamW steps, round-trips through a file and compares model and optimizer states exactly. Then extend it in your own repository with scaler, RNG and DataLoader order plus a resume-equivalence check.
Self-test
-
The checkpoint holds model, optimizer and iteration. Why is the iteration number needed?
It restores the schedule position. Without it the learning rate restarts at the wrong phase, for example back at the warmup peak, and the run diverges from the original. -
Your resume test matches the uninterrupted run exactly for five steps on CPU, but on a GPU the losses drift apart by about 1e-6 per step. Is resume broken?
Not necessarily. Nondeterministic kernels and reduction order give tiny float differences that compound, so state the equivalence level: exact on CPU with deterministic algorithms, a tolerance band on GPU. A sudden jump at the resume step, by contrast, means missing state. -
Resumed loss diverges from the uninterrupted run within a few steps. Name two suspects.
Unrestored RNG orDataLoaderorder, which feeds different batches, and missing scaler or optimizer moment state, which changes the update itself. -
Why pin
tokenizer_configandconfigversions alongside the checkpoint?Token IDs and tensor shapes are config-relative. A changed vocabulary or width silently redefines what every saved number means, so the binding must fail loudly on mismatch.
Completion criterion
Stop when
Your checkpoint round-trips through test_checkpointing, and your own resume test matches the uninterrupted loss to your stated tolerance.
Leads to
- 3.6 Decoding policies: from greedy to nucleus sampling Generate from resumed weights with owned sampling policies.
- 3.7 Stopping rules and reproducible generation Extend determinism from training into generation.
Comes after
- 1.4 A tokenizer you can trust: round trips, identity and streamingMake tokenizer compatibility part of checkpoint loading.
- 1.7 Filtering, mixtures and deterministic shardsResume the data stream alongside optimizer and model state.
- 3.3 AdamW, schedules and the global token batchPersist optimizer moments and the schedule position with the model.
- 3.4 Keeping training numerically healthyAdd the scaler state to the checkpoint so protection survives a restart.
Resources for this unit
- LLMs from Scratch companion - training and generation Use the ch05 training and generation code for the small pretraining loop. Add full checkpoint state, explicit loss masks and resume-equivalence tests that the base example does not supply.
- Stanford CS336 Assignment 1 - training and decoding contract Use cross_entropy, adamw, adamw_accounting, learning_rate_schedule, gradient_clipping, data_loading, checkpointing, training_together and decoding specifications plus public tests. Debug on CPU and implement independently in your own repository.
- Modern LLM Notebook - configuration, loss and generation Use English notebooks 09-model-config, 10-training-loss and 20-generation to visualize configuration ownership, shifted loss and decoding policies. Change edge cases and cross-check probability calculations against your own implementation.
- PyTorch saving, loading and reproducibility references Use the state-dict and reproducibility notes to enumerate all state needed for the supported resume guarantee. Test model, optimizer, schedule, scaler, data order and RNG state across an interruption.
Also in this section
- Hugging Face LLM Course - library translation Only after the blank-file loop passes, map tokenization, model configuration, data collation, generation configuration and stopping behavior into Transformers APIs. Keep the from-scratch tests as the authority.
- PyTorch Automatic Mixed Precision recipe Use the
autocast,GradScaler, unscale-before-clip, scaler checkpointing and inf-NaN troubleshooting sections to add mixed precision to the blank-file loop. Keep the CS336 clipping contract as the authority for the clipping rule.