Save Enough State to Repeat or Resume the Run

Separate inference reload, a comparable new run, and exact training continuation. Save model, optimizer, step, configuration, vocabulary, data identity, environment, and random-generator state accordingly.

Saving trained weights is enough to reproduce inference under the same model definition. Continuing training exactly requires more state because AdamW and the batch sampler remember the past.

Separate Three Reproducibility Goals

Goal	Required state
reproduce current logits	model definition, configuration, vocabulary, model weights, evaluation mode
start a comparable new run	code and environment, configuration, corpus identity, split and model seeds
continue the interrupted run	all above plus optimizer state, step, scheduler if any, and random-generator states

The frozen checkpoint record contains:

Every entry has a reason. AdamW stores moving averages and a step count. The batch generator determines the next examples. Vocabulary order determines what embedding row and logit column each ID means.

Test Loading, Not Merely Saving

A successful file write proves little. Construct a new model and optimizer, load both state dictionaries, restore the generator, and verify:

logits before the next update match exactly in the same environment;
the next batch indices match;
the next loss and gradients match;
the next updated parameters match.

The verified CPU experiment passes this one-step continuation test.

Seeds Limit Variation; They Do Not Guarantee Portability

PyTorch warns that complete reproducibility is not guaranteed across releases, platforms, or CPU and GPU execution. Record at least Python version, PyTorch version, device, operating system, and relevant package versions with the run.

For this CPU toy model, separate generators make model initialization, corpus sampling, and batch order independently repeatable. Larger data pipelines may also require worker and library seeds.

Q1. Choose state for exact continuation

You saved only model.state_dict() after step 500. Can you guarantee the same step 501 that the uninterrupted AdamW run would have taken?

Choose one

No. Optimizer and random-generator state are missing, so the moments and next batch may differ.

Yes. Model weights uniquely determine every future optimizer update.

Yes, if the model seed is written in the filename; optimizer state is unnecessary.

Select one choice, then check.

Hint

AdamW stores moments, and sampling has its own state.

Solution

No. Restore optimizer state, step and any scheduler state, and the random-generator states that determine the next batch, in addition to model weights and configuration.

References

These are current framework instructions and should be rechecked when the project changes its supported PyTorch release.