Neural networks, deep learning, NLP · insight-1403
Checkpoint the optimizer, scheduler and step counter, not only the weights
- Evidence type
- documented fact
- Difficulty
- easy
- Roles
- ML engineer
- Industries
- cross-industry
Summary: situation, action, reason, boundary
Situation
A long training job is preempted and you resume it from a file that holds only the model weights.
Action
Save one dictionary with model, optimizer and scheduler state, the global step, the scaler state if used, and the random-number-generator states. After loading, call the training-mode switch before continuing.
Reason
Optimizers keep buffers that change as training proceeds, so the model's weights alone are not enough to continue the same optimisation trajectory, as the PyTorch checkpoint guidance states.
Boundary
Any run longer than a few minutes, and every spot or preemptible instance.
Symptom or risk
Adam-style moment estimates and the schedule position restart from zero. The loss jumps after resuming, and a run that was interrupted no longer matches the uninterrupted run you wanted to reproduce.
Exceptions
A final inference checkpoint needs only weights; keep it as a separate artefact.
Trade-offs
Optimizer state roughly doubles to triples checkpoint size for Adam-type methods.
Pitfalls
Loading state into an optimizer built over parameters in a different order maps buffers to the wrong tensors.
Example
Hypothetical: resuming a booking-volume model from weights alone shows a loss spike for a few hundred steps as the moment estimates are rebuilt; resuming from the full dictionary does not.
Lessons and concepts
Claims and their sources
The PyTorch guidance for resuming training saves the optimizer state_dict because it holds buffers and parameters that are updated as the model trains.