Saving and Loading PyTorch Model Weights
A saved model can be resumed: load the model and optimizer in a new Python session or notebook, then continue pretraining for one more epoch with train_model_simple.
Try it: while Loops, break and continue
How a while loop re-evaluates its condition before every pass, how break leaves the loop at once, how continue jumps back to the condition, how a sentinel value ends a while True loop, and why a loop whose condition never becomes False runs forever.
How it works
- Evaluate the condition. If it is False, skip the body and continue after the loop.
- If it is True, run the body from the top, then go back and evaluate the condition again.
- continue skips the rest of the body and goes straight back to the condition check.
- break leaves the loop immediately; the condition is not checked again (while True loops rely on it, e.g. stopping at a sentinel like 'done').
- If nothing in the body can make the condition False, the loop never ends; the lab's guard stops it after N iterations and says so.
Default run (21 steps): Ready: Countdown. A while loop checks its condition before every pass. … Finished after 5 iterations and 6 condition checks.
Simplified: The lab's own small Python interpreter runs a fixed, hand-written script (you choose the variant and the values, never the code); nothing is executed as real Python. Every script was run in CPython 3.12 with the same inputs and the lab reproduces its output and error messages. input() returns the lines you type in the list, one per call (the prompt is not printed), and raises EOFError when they run out, like a closed keyboard stream. The guard is the lab's, not Python's.
Loading the simulation…
Why Training Must Survive a Break
Training a large language model takes a long time, and a training run does not always finish in one sitting. A session may end, a notebook may be closed, or a machine may need to be restarted. If the work done so far were lost each time, every interruption would force the model to begin again from nothing. Saving and loading model weights is the mechanism that prevents this. Instead of treating each training session as a fresh start, the progress made so far can be stored and then picked up again later.
A useful way to think about a saved model is as a snapshot of a training run. A saved checkpoint is not a finished product by default. It captures where training stood at one moment, and it can be restored later and pushed further along. The central question of this article is simple: when you stop training and come back later, what exactly has to survive the break, and what has to be restored before training can carry on from where it left off?
Following the Weights Across a Break
Follow the weights through one full cycle. First, they exist inside a model during a training session, where their values are updated as training proceeds. Second, the current values are written to storage. Third, in a later session, the saved values are read back in and placed into a model, so that the model starts from that stored state rather than from scratch. Fourth, the optimizer is restored as well, and training continues for one more epoch with the train_model_simple function.
Read this flow from left to right, because each stage depends on the one before it. The saved weights have no effect until they are loaded back into a model. Training can only continue once both the model and the optimizer are in place. The final stage uses the same train_model_simple call that was used for the original training, now applied for one additional epoch.
What Counts as Weights
Model weights are the weight parameters stored in the .weight attributes of PyTorch Linear and Embedding layers. These are the values that save and load operations carry from one session to the next. When a model is said to have been saved, the thing being saved is the collection of these values.
Saving model weights stores the learned values of a model so that they can be used later. Loading model weights reads previously saved weights back in, so the model starts from that saved state.
The word weights has a specific meaning in this chapter. It refers to the learned values held by Linear and Embedding layers, not to every piece of information a program might keep about a model.
Locating Weights with model.parameters()
A GPTModel is built from smaller parts. It contains Linear and Embedding layers, and each of those layers exposes its weights through its .weight attribute. The layers sit inside the model in a nested structure, so the weights are at the bottom of a tree of components rather than sitting directly on the model object.
You reach these parameters through model.parameters(). This is the same method that was used earlier in the chapter to inspect the model during training. Walking the model's structure with model.parameters() is how the weight values become visible, and they are the values a save operation needs to capture.
What a Resumable Save Contains
It is tempting to think that saving the model is enough, because the model is the part that makes predictions. The source is more specific than that. When training is continued, both the model and the optimizer are loaded. The model carries the learned weight values, while the optimizer is the component that drives the update process during training. Restoring both lets training continue as a single run, rather than as a fresh start with a partly trained model.
Resuming One More Epoch
The source describes the resumed workflow directly. After the weights are saved, you load the model and optimizer again, possibly in a completely new Python session or in a Jupyter notebook file, and then continue pretraining for one more epoch using the train_model_simple function. The function is the same one used for the original training, so the resumed run follows the same training procedure as before.
Continuing a training run in a new session
A model was trained in one Python session and its progress was saved. A new session must pick up that progress and train for one more epoch without starting over.
- 1Save the trained state
Write the current weight values of the model to storage, so the progress from the first session survives the end of that session.
- 2Create a fresh model in the new session
The new session starts with an empty model. Its weights are not yet the trained values, so it has not yet recovered any progress.
- 3Load the saved weights into the model
Read the previously saved weights back in, so the model starts from the stored state rather than from scratch.
- 4Load the optimizer
Restore the optimizer as well, since it drives the update process. Together with the model, this lets training continue as one run.
- 5Run train_model_simple for one epoch
Call the same training function for one additional epoch. The weights are updated from the restored state, not from fresh random values.
Training resumes from the saved state and continues for one more epoch, so the second session extends the first session's progress instead of repeating it.
Starting from Pretrained GPT-2 Weights
The earlier training in this chapter used a small GPT-2 style model on a limited dataset, a short story book. That kept the focus on fundamentals, but the model had seen very little text. OpenAI openly shared the weights of its GPT-2 models. Using those shared weights means you do not have to retrain a model on a large corpus yourself, which would demand a very large investment of money and computing resources.
The plan is to load these pretrained weights into the GPTModel class built earlier, and then use that model for text generation. The same pretrained weights are reused in chapter 6 to fine-tune the model for a text classification task and to follow instructions in a way similar to ChatGPT. The loading step follows the same idea as any other load: the saved weight values are read in and placed into the model's weight parameters, so the model begins from that stored state.
Mistakes That Make a Run Start Over
Common mistake Assuming that saving the model is enough to resume training
Why it failsThe model carries the learned weight values, but the optimizer drives the update process. Without it, the run is not restored as a single continuing run.
Do this insteadSave and load both the model and the optimizer when you intend to continue training.
Common mistake Loading only the weights into a fresh model and expecting the run to continue
Why it failsRestoring only part of the state leaves the training process incomplete, so the result is a fresh start with a partly trained model rather than one continuous run.
Do this insteadRestore the model state and the optimizer state together before calling train_model_simple.
Common mistake Creating a new model and skipping the load step
Why it failsA newly constructed model begins without the stored progress, so training from it repeats work that was already done.
Do this insteadLoad the saved state into the model before continuing training.
Common mistake Treating a save as if it were a load
Why it failsA save without a later load stores the progress but does not put it back to work.
Do this insteadRemember that saving and loading are a pair. The load is what brings the saved state back into training.
Common mistake Treating a saved checkpoint as a finished product
Why it failsA saved checkpoint is a snapshot of a training run that can be pushed further along, not a result that ends training.
Do this insteadContinue training from the restored state when more epochs are needed.
| Setup | What is restored | How training behaves |
|---|---|---|
| Model and optimizer both loaded | Learned weight values and the optimizer state | Training continues as a single run from the saved snapshot |
| Only the model weights loaded | Learned weight values, without the optimizer | Training does not carry on as a single run; it behaves like a fresh start with a partly trained model |
| Fresh model created, nothing loaded | Nothing from the earlier session | Training starts over from the model's initial state |
How the restored components determine whether a resumed run continues or restarts
Predict Before You Reveal
What do you think happens?
If you load only the model weights into a new session and skip the optimizer, does training continue as a single run?
Reveal answer
No, the run behaves like a fresh start with a partly trained model.
The model carries the learned weight values, but the optimizer drives the update process. Restoring both lets training continue as a single run rather than as a fresh start with a partly trained model.
Describe, in your own words, the order of operations for resuming training in a new Python session. Name each step from the moment the weights are saved to the moment train_model_simple runs for one more epoch, and say which component must be restored alongside the model.
Hints
- Start from the save step and work forward, one stage at a time.
- Remember that the optimizer is restored together with the model.
- The final call is the same
train_model_simplefunction, applied for one additional epoch.
Key Takeaways
- Saving model weights stores the learned values of a model, and loading reads those values back in so that the model starts from the saved state rather than from scratch.
- Model weights are the parameter values in the .weight attributes of Linear and Embedding layers, and they can be reached through
model.parameters(). - Resuming training requires restoring both the model and the optimizer, because the optimizer drives the update process and restoring both lets training continue as one run.
- A resumed run can be carried on with
train_model_simplefor one more epoch, and pretrained GPT-2 weights shared by OpenAI can be loaded into a GPTModel instead of retraining on a large corpus.