Click any words you think contain an error. Click again to unmark.
A faithful resume needs weights, optimizer moments, data loader position, and RNG state; this plan restores only weights and oversells the result.
Picture pausing a long board game and coming back the next day. The pieces on the board are the model weights. The little notes you wrote about whose turn it is, what cards each player saw, and how shuffled the deck got are the optimizer moments and the data pipeline state. If you put the pieces back but throw away the notes and reshuffle from scratch, you are not resuming the game. You are starting a similar new game. The plan above keeps the pieces and forgets the notes, then claims it is the same game. That is the gap.
Concept explanation~2 min read
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
Concept explanation~2 min read
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
A checkpoint feels like a snapshot of a model, but during pretraining it is a snapshot of a process. The process has three coupled pieces of state, and a resume that restores only one of them is not a resume at all.
Why optimizer state matters as much as weights
AdamW keeps two running quantities per parameter: the first moment (an EMA of gradients) and the second moment (an EMA of squared gradients). The update is roughly:
At step 90000 those moments hold a calibrated estimate of each parameter's recent gradient mean and scale. The denominator √v̂_t is what normalizes the effective step size per parameter. Throw it away and every parameter starts the next step with the same nominal scale, regardless of how steeply or gently its loss surface moves.
Now stack that with a mid-run learning rate. A WSD schedule holds peak η for most of the run, and a cosine schedule stays close to peak for a long stretch after warmup. Combining fresh moments with peak LR applies full-size updates before the optimizer has any estimate of each parameter's gradient scale. The result is a few hundred steps of huge, badly directed updates, often visible as a sharp loss spike, sometimes as outright divergence. This is a common cause of bad resumes.
Situations where this technique stops working.
2–4 min · Everything important, quickly.
Real products, models, and research that use this idea.
- DeepSpeed's save_checkpoint stores model, optimizer, and LR scheduler state and takes a client state dictionary for the training step, so the loader can be fast-forwarded to the right batch after a node failure; Megatron-LM records the consumed-sample count for the same purpose
- Meta's Llama 3 report counts 466 job interruptions during a 54-day period of 405B pretraining, so restoring the full training state after a failure is routine work at that scale
What an interviewer would ask next. Try answering before peeking at the approach.
QHow would you debug a resumed run that diverges within a few hundred steps even though you restored all three legs?
Check determinism settings (cuDNN, NCCL, attention kernel), verify the LR schedule step matches, confirm gradient accumulation and global batch size match, and inspect whether mixed-precision scale factors were also restored.
Red flags & common mistakes
The phrases that signal junior thinking. Click to expand.
Red flags & common mistakes
The phrases that signal junior thinking. Click to expand.
Treating a checkpoint as weights only and assuming Adam moments will rebuild safely at mid-run learning rate, then blaming the resulting loss spike on a bad batch.
60 second bullets to scan on the way to the call.
What state does a training checkpoint need to contain for a faithful resume?
Why do fresh AdamW moments at peak learning rate cause divergence?
Primary sources. Browse if you want the original framing.
Same topic, related formats. Practice these next.