A checkpoint is more than its state dict

By Ajish Pradeep · 2026-10-07 · 6 min read

Claim. A model or a training run can load without an error and still differ from the one that was saved. In my tests a normalisation constant missing from the state dict changed outputs by up to 30, and a restored optimizer silently kept the old learning rate and rewrote the checkpoint it came from.

"Save the state dict, load the state dict" is how most of us think about checkpoints, and when the load succeeds without an error we assume we have the same model back. I stopped assuming that while building an exact replay of a training step for another experiment, where two loads that should have been identical weren't. That made me want to list what a state dict doesn't carry, and check each item on purpose instead of meeting it by accident. Each of these four loaded without an error:

What was restored	What differed afterwards
Weights into a batch-norm layer with another $\epsilon$	outputs, by up to 30.1 on values of about 4.2
Weights through a loader that keeps matching shapes	the resized output layer, skipped and left random
An optimizer, constructed with a new learning rate	the rate, back to the saved one; and the checkpoint itself, rewritten by later steps
An optimizer without its scheduler, mid-run	the learning rate 100 steps later, 41% higher than planned

A constant that isn't in the state dict

A batch-normalisation layer stores its scale, its shift and its running mean and variance as tensors, and all of those go into the state dict. The small constant $\epsilon$ that keeps the division stable does not. It is a constructor argument:

y = \gamma\, \frac{x - \hat\mu}{\sqrt{\hat\sigma^2 + \epsilon}} + \beta

Load the weights into a layer built with a different $\epsilon$ in PyTorch (Paszke et al., 2019¹) and nothing complains. Whether the outputs change depends on how $\epsilon$ compares with the stored variances. For channels with a healthy variance it changes almost nothing. For a channel whose variance is close to zero, which trained networks do have, the denominator is set by $\epsilon$ itself.

I built a small network with one batch-norm layer and gave it running variances spread from $10^{-5}$ to 1, to stand in for a trained network with a few nearly constant channels. Then I loaded its state dict into the same architecture built with $\epsilon = 10^{-5}$ instead of $10^{-3}$ , two values that are both common defaults. The load reported no problem. The largest output difference was 30.1, on outputs whose largest magnitude was 4.2. The keys in the state dict for that layer were the weight, the bias, the running mean, the running variance and a batch counter, and nothing that records $\epsilon$ .

The fix is to rebuild the module from the same configuration, or to carry the configuration in the checkpoint and check it on load. Copying weights into a freshly constructed module is safe only when every constructor argument matches.

A head that was skipped

When an output layer changes size, for example to predict more classes, a common way to reuse a checkpoint is a loader that keeps only the tensors whose name and shape both match. In my test the checkpoint held four tensors and the new model expected an output layer of a different size. The loader kept two and skipped the output layer's weight and bias without a word, so the new head started from random values. PyTorch's own loader is stricter here: even in its non-strict mode it raised an error on the shape mismatch. The silent version is a pattern people write around it.

Its danger depends on what is measured afterwards. Any metric that doesn't depend on the skipped tensors looks normal, so a run can train for a long time on a head that started from scratch while the shared parts of the model report healthy numbers. The cheap guard is to log the number of tensors loaded against the number in the checkpoint, and to fail when they differ in a way nobody planned.

An optimizer that shares and overrides

Restoring an optimizer surprised me twice. I created an Adam optimizer with a new learning rate of $10^{-3}$ and loaded a saved optimizer state into it. Afterwards its learning rate was $10^{-2}$ , the rate stored in the checkpoint. The saved parameter groups, with their learning rates, replace the ones the optimizer was constructed with. To change the rate on resume, it has to be set after loading.

Then I trained the restored model for three steps and looked at the checkpoint I had loaded from, which was still in memory. Its stored first-moment estimate had changed by up to 0.197. Loading did not copy the moment tensors, because they were already on the right device with the right type, so the restored optimizer and the checkpoint shared memory, and every step wrote into both. One resume from an in-memory checkpoint is fine; a second one, from the same object, starts from a state the first resume has already advanced. A deep copy before loading avoids it. This was on PyTorch 2.4; other versions may copy.

A schedule that restarts

The last item is the most familiar and still easy to miss. A cosine schedule over 1,000 steps, resumed at step 500 with the optimizer restored but not the scheduler, starts a new cosine curve from the current rate. At step 600 the uninterrupted run had a learning rate of 0.0345. The run resumed with the scheduler's state saved and restored matched it exactly. The run resumed without it had 0.0488, 41% higher, and would keep diverging from the plan for the rest of training. The same applies to anything else with internal state: warmup counters, loss scalers for mixed precision, exponential moving averages of the weights, and the random number generators behind data order and augmentation.

What a complete checkpoint holds

From these checks, the list I would now use:

the model's state dict and the configuration used to construct it, including constants such as $\epsilon$ and momentum;
the optimizer's state, deep-copied, with any change of learning rate applied after loading;
the scheduler, any loss scaler and any weight average;
every random number generator's state, and the position in the data order;
the number of tensors loaded, checked against the number saved;
a fingerprint test: one fixed batch through the model before saving and after loading, with the outputs compared exactly.

The last check is the one that would have caught every problem above in a single line.

The exact replay that surfaced the optimizer traps is in taking a gradient spike apart, and what a fresh optimizer does to a trained model when its state is lost is in restarting Adam moves every weight by the full learning rate.

What I still don't know

I tested one framework version, and the optimizer behaviour in particular depends on it. I also haven't checked how these issues carry over to exported formats, where the constructor configuration is baked into a graph and a mismatch shows up differently, or to models split across devices, where loading has more ways to place a tensor somewhere unexpected.

References

PyTorch: An Imperative Style, High-Performance Deep Learning Library — Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang, Zach DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai and Soumith Chintala, NeurIPS 2019