Loading Model Weights and Continuing Training
Adaptive optimizers such as AdamW store additional parameters for each model weight, and that data is separate from the weights themselves.
The Belief That Weights Are the Whole Model
Many learners treat a saved model as a single object: the numbers that make up the weights. If those numbers are written to disk and read back later, the reasoning goes, training can pick up exactly where it stopped. This belief feels natural, because the weights are what the model learned. It is also incomplete, and the gap matters most when the optimizer is an adaptive one such as AdamW.
This article is built around that belief. You will first see what goes wrong when the optimizer is left out of a checkpoint, then follow the code that saves and restores both pieces, and finally use the same reasoning to find the point where a resumed run stops matching the original.
What do you think happens?
You save a model and its AdamW optimizer, then later create a fresh model and a fresh AdamW optimizer, load the model weights, and continue training. Does the new optimizer behave as though it has seen the earlier training history?
Reveal answer
No. The fresh optimizer does not carry the history the earlier optimizer built up, so it does not behave as though it saw that training.
The weights were restored, but the optimizer state that AdamW keeps alongside them was not. Without restoring that state, the optimizer resets on restart.
What AdamW Keeps Besides the Weights
A plain gradient-descent style optimizer applies a fixed rule to every update. Adaptive optimizers such as AdamW do not work that way. They use historical data to adjust the learning rate for each model parameter dynamically. To make that possible, AdamW stores additional parameters for each model weight.
Those additional parameters belong to the optimizer, not to the model's weights. This is the central point of the whole article. The model's state_dict describes what the network has learned. The optimizer's state describes the running history that the optimizer uses to decide how to change those weights next. Two objects, two different kinds of data, and only one of them is usually thought of as 'the model'.
The diagram below shows the layout for one small weight tensor. The model holds the weight values that the forward pass uses. The optimizer holds a separate set of per-weight estimates, and each estimate is tied to one particular weight position. Restoring the weights alone leaves the second group empty.
Notice that the optimizer's entries are not part of the forward computation at all. The network produces its outputs from the weights alone. The optimizer's entries only matter when the next update is calculated, which is why a model can load and run without error even when the optimizer state is missing.
What Happens When the Optimizer Resets
The source is direct about the consequence. If the optimizer state is not saved, the optimizer resets on restart. Its accumulated history is gone, and it begins again as if no training had happened before. Because the weights were preserved, the reset is easy to miss, and the run appears to continue normally.
The effect on learning can be serious. The source notes that the model may then learn suboptimally, or even fail to converge properly. It describes that failure in strong terms, noting that the model may lose the ability to generate coherent text. The misconception this article targets is the assumption that a reset optimizer is harmless because the weights were preserved.
Saving Both Pieces Into One File
PyTorch lets you save the model and the optimizer together. The model contributes its state_dict, and the optimizer contributes its own state_dict. Both are placed in one dictionary under two keys, and that dictionary is written to a single file with torch.save. The key names used below, model and optimizer, are a choice made for this example.
build_model()stands in for whatever code defines your network- The optimizer is created from
model.parameters(), so it exists only after the model does - Both state dictionaries go into one dictionary under two keys
- torch.save writes that one dictionary to a single file
Notice what this call does and does not do. It does not save the training loop, the data loader, or the code that defines the model. It saves the two state dictionaries, and only those, into the file named in the call. A later restore depends entirely on what this one file contains, so anything missing from the dictionary is missing from the resumed run.
Restoring in the Correct Order
Restoring is a two-stage operation. First the saved file is read back into memory with torch.load. Then the pieces are applied to freshly created objects with load_state_dict. The file does not reconnect itself to any live model or optimizer. That connection is made by you, in code, and the order of the steps matters.
The sequence below follows the restore path. The script reads the file first, then creates the model, then creates the optimizer on the model's parameters, and only then applies each saved state. Each load_state_dict call depends on a matching object already being in place. The optimizer is created after the model's parameters exist because its constructor needs model.parameters(). Its saved state is applied only after it exists.
- torch.load reads the file into memory as a dictionary
- The model is built before anything is loaded into it
- The optimizer is built after the model, because it needs
model.parameters() - Each
load_state_dictcall receives the matching entry from the dictionary
Tracing one restore
A resume script loads the model correctly but skips the optimizer restore. The developer asks where the resumed run first departs from the original run.
- 1Read the file
torch.load returns the dictionary with the model and optimizer entries. This step succeeds.
- 2Build the model
The model object exists, so model.load_state_dict has a target to fill.
- 3Build the optimizer
A new AdamW optimizer is created with its own empty history. It has no saved per-weight entries yet.
- 4Load the model weights
model.load_state_dict(ckpt['model'])places the saved weights into the model. The weights now match the original run. - 5Skip the optimizer restore
optimizer.load_state_dict(ckpt['optimizer'])never runs. The optimizer still holds the reset state from its constructor. - 6Take the first update
The update uses the reset history instead of the saved one, so it differs from the update the original run would have made.
The resumed run diverges at the first update after restart. The model weights loaded without error, which is why the fault is easy to miss. The single line most likely responsible is the missing optimizer.load_state_dict call, because it is the only step that gives the optimizer its earlier history.
Inspecting What a Checkpoint Holds
When a restore does not behave as expected, it helps to look inside the file before loading it into objects. The check below uses the same key names as the save example. It shows which entries exist and which load call each entry feeds. It does not reproduce the training code.
['model', 'optimizer']
<class 'dict'>
<class 'dict'>keys()lists the top-level entries, which should show both expected names- Each entry is inspected before it is passed to a
load_state_dictcall
If the key list is missing either name, the corresponding restore call has nothing correct to read. A check like this separates a problem in the file from a problem in the restore code, which is the first distinction to make when a resumed run behaves oddly.
Mistakes That Break a Resumed Run
Common mistake Saving only the model state_dict
What it looks liketorch.save(model.state_dict(), 'checkpoint.pt')Why it failsThe optimizer's per-weight history is not in the file, so the optimizer resets on restart.
Do this insteadPut both state_dicts into one dictionary under two keys and save that dictionary.
Common mistake Calling optimizer.load_state_dict before the optimizer exists
Why it failsEach
load_state_dictcall depends on a matching object already being in place. The optimizer has to be constructed first, frommodel.parameters().Do this insteadBuild the model, then build the optimizer, then apply both saved states.
Common mistake Assuming a successful model load means the run resumed correctly
Why it failsThe weights can load without error while the optimizer starts from reset history, so the next updates differ from the original run.
Do this insteadCheck that the optimizer restore runs and that its entry comes from the optimizer key.
Common mistake Treating the checkpoint as a record of the training loop, when it stores only the two state dictionaries
Common mistake Reading the file with torch.load but never applying its entries to newly created objects
Finding the Point of Divergence
The restore path has a fixed order: read the file, build the model, build the optimizer, apply the model state, apply the optimizer state. A resumed run can stop matching the original run at any point in that path where a required step is skipped, reordered, or applied to the wrong object.
A resume script runs torch.load, builds the model, calls model.load_state_dict(ckpt['model']), and then builds a fresh AdamW optimizer and starts training. Each loss value is close to the original run's value only for the first few steps, and then the updates drift away. Identify the step where the resumed run diverges, and explain in terms of what the optimizer is missing.
Hints
- Compare the script's steps with the six-step restore order in the sequence diagram.
- Ask what the optimizer's entries contain at the moment the first update is computed.
- Consider whether the saved optimizer entry is ever applied to the new optimizer object.
Key Takeaways
- An adaptive optimizer such as AdamW stores additional per-weight parameters that are separate from the model's weights, so the weights alone do not capture its training state.
- If the optimizer state is not saved, the optimizer resets on restart, which can leave the model learning suboptimally or failing to converge even though the weights loaded correctly.
- torch.save can place the model's
state_dictand the optimizer'sstate_dictinto one dictionary under two keys and write that dictionary to a single file. - A restore reads the file with torch.load, builds the model, builds the optimizer from
model.parameters(), and then applies each part withload_state_dictin that order. - When a resumed run diverges, check the restore sequence step by step, starting with whether the optimizer's saved state was actually applied to the optimizer.