Training a GPT Model with a Simple Training Loop

Fine-tuning a saved GPT model starts by restoring two pieces of state: the model and the optimizer, both stored in one checkpoint file.

Programming

Picking Up Where Training Stopped

Training a GPT model is expensive, so you rarely want to throw away work that has already been done. Fine-tuning a pretrained model means picking up from a saved point and continuing training for more passes over the data. This article follows that resume workflow for the GPT model built earlier in the course. The goal is to see exactly what has to be restored so that continued training starts from the saved point rather than from scratch.

The source describes the whole job in one sentence: load the model and the optimizer that were saved in the main chapter. The rest of this article unpacks what that sentence requires, one piece at a time, and explains why the order of the steps matters.

What do you think happens?

A saved checkpoint is loaded into a single dictionary with two entries you care about. Which entry should go to model.load_state_dict(), and which should go to optimizer.load_state_dict()?

model_state_dict goes to the model, optimizer_state_dict goes to the optimizer
optimizer_state_dict goes to the model, model_state_dict goes to the optimizer
Both entries go to the model, because the optimizer is rebuilt from it
Only the model entry is needed, because the optimizer restarts on its own

Reveal answer

Answer:

model_state_dict feeds model.load_state_dict(), and optimizer_state_dict feeds optimizer.load_state_dict().

The two keys are the contract between saving and loading. Each saved entry holds the state of one specific object, so each one must reach the matching load call. One call cannot serve both, because the model and the optimizer hold different state.

Two Pieces of State in One File

A saved checkpoint holds two separate pieces of state. The first is the state of the GPT model, which is its learned parameters. The second is the state of the optimizer, which in this workflow is torch.optim.AdamW. Both live in the same checkpoint file, but they are different things and each must be restored into its own object.

The optimizer keeps its own internal state alongside the model. This is why the optimizer must be saved and restored too. If you restored only the model, the optimizer would behave as though it had never run, even though the model's parameters had already been trained.

How the model and the optimizer relate to each other and to the single checkpoint file, and what each one holds.

Read the diagram as two arrows leaving one file. The file is the shared source, the model and optimizer are two separate receivers, and the optimizer is attached to the model's parameters. That attachment is why the optimizer has to be built after the model exists.

Reading the Checkpoint by Key

The loading code reads the checkpoint into a single object with torch.load(). It then pulls two entries out of that object by key. The key model_state_dict belongs to the model, and the key optimizer_state_dict belongs to the optimizer. These two keys are the whole contract between saving and loading, so they are the first place to look when something goes wrong.

python

import torch

# checkpoint_path is a placeholder for the file saved in the main chapter
checkpoint = torch.load(checkpoint_path)

print(list(checkpoint.keys()))   # the two entries you care about are listed here

model_entry = checkpoint["model_state_dict"]
optimizer_entry = checkpoint["optimizer_state_dict"]

import torch

# checkpoint_path is a placeholder for the file saved in the main chapter
checkpoint = torch.load(checkpoint_path)

print(list(checkpoint.keys()))   # the two entries you care about are listed here

model_entry = checkpoint["model_state_dict"]
optimizer_entry = checkpoint["optimizer_state_dict"]

Output (expected)

['model_state_dict', 'optimizer_state_dict']

Notes on the code

torch.load() reads the file into one dictionary object
Each entry is retrieved by its key, not by position
The keys are the contract between the code that saved the file and the code that loads it

Which key in the checkpoint file feeds which loading call, and where a mismatch between saved and rebuilt objects would show up.

Building Before Loading

The saved state is a set of values, not a running model. Something has to exist to receive those values. The loading sequence therefore creates a fresh GPTModel using the configuration named GPT_CONFIG_124M, and only then calls model.load_state_dict() with the model_state_dict entry. Creating the object first gives the loading call a structure to fill in.

The optimizer is handled the same way. It is constructed first, then given its saved state. In the source it is created as torch.optim.AdamW, applied to the parameters of the freshly loaded model, with a learning rate of 5e-4 and a weight decay of 0.1. Only after that is optimizer.load_state_dict() called with the optimizer_state_dict entry.

Which objects are created first, which saved states are injected into them, and why the objects must exist before their states are loaded.

The sequence is fixed: read the checkpoint, rebuild and restore the model, rebuild and restore the optimizer, and then train. Each load call needs its receiving object to exist already, which is why the construction steps come before the load steps in every branch of the flow.

Tracing a Full Resume

python

import torch

checkpoint = torch.load(checkpoint_path)   # placeholder path to the saved file

model = GPTModel(GPT_CONFIG_124M)         # 1. build the structure
model.load_state_dict(checkpoint["model_state_dict"])   # 2. fill it

optimizer = torch.optim.AdamW(
    model.parameters(), lr=5e-4, weight_decay=0.1
)                                        # 3. build the optimizer on the restored model
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])   # 4. fill it

# Continue training from the restored state for one more epoch.
# Other arguments are as used in the main chapter.
train_simple_function(..., num_epochs=1)

import torch

checkpoint = torch.load(checkpoint_path)   # placeholder path to the saved file

model = GPTModel(GPT_CONFIG_124M)         # 1. build the structure
model.load_state_dict(checkpoint["model_state_dict"])   # 2. fill it

optimizer = torch.optim.AdamW(
    model.parameters(), lr=5e-4, weight_decay=0.1
)                                        # 3. build the optimizer on the restored model
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])   # 4. fill it

# Continue training from the restored state for one more epoch.
# Other arguments are as used in the main chapter.
train_simple_function(..., num_epochs=1)

Notes on the code

The model is created before its state is loaded
The optimizer is created on model.parameters() after the model is restored
num_epochs=1 trains for one more epoch from the restored state

Worked example

Resuming for one more epoch

You have a checkpoint file saved from the main chapter. Describe the steps needed to continue training the GPT model for one more epoch.

1
Read the checkpoint
Use torch.load() to read the file into one dictionary with the entries model_state_dict and optimizer_state_dict.
2
Rebuild and restore the model
Create a fresh GPTModel from GPT_CONFIG_124M, then call model.load_state_dict() with the model_state_dict entry.
3
Rebuild and restore the optimizer
Create torch.optim.AdamW on the parameters of the restored model with learning rate 5e-4 and weight decay 0.1, then call optimizer.load_state_dict() with the optimizer_state_dict entry.
4
Train for one epoch
Call train_simple_function with num_epochs=1. Training begins from the restored state, not from random initial values.

Result

The model and optimizer both resume from the saved point, and one additional epoch of training is run on top of that state.

Mistakes When Resuming

Common mistake
Restoring only the model and skipping the optimizer
Why it fails
The optimizer keeps its own internal state alongside the model, so without it the optimizer behaves as though it had never run.
Do this instead
Load both model_state_dict and optimizer_state_dict from the same checkpoint.
Common mistake
Calling load_state_dict before the object exists
Why it fails
The saved state is a set of values and needs a structure to fill. The load call has nothing to receive the values if the object has not been created.
Do this instead
Construct GPTModel from GPT_CONFIG_124M first, and construct the optimizer on the restored model's parameters first, then load each state.
Common mistake
Swapping the two checkpoint entries
What it looks like
```
optimizer.load_state_dict(checkpoint["model_state_dict"])
```
Why it fails
Each entry holds the state of one specific object, so the wrong entry will not match the receiving object.
Do this instead
```
Send model_state_dict to model.load_state_dict() and optimizer_state_dict to optimizer.load_state_dict().
```
Common mistake
Assuming the optimizer can be built before the model
Why it fails
The optimizer is applied to the parameters of the freshly loaded model, so the model has to exist first.

Running One More Epoch

EASY

Suppose a teammate wrote the resume code with the optimizer entry passed to the model's loading call, as model.load_state_dict(checkpoint["optimizer_state_dict"]). Trace the code: which load call is affected, and where would you expect the problem to show up? Then state the single call that continues training for one more epoch from the restored state.

Hints

Compare each checkpoint key with the object its state belongs to.
Look at which load call receives the wrong entry, not at the torch.load() line.
The continued-training call is train_simple_function with num_epochs=1.

Your attempt (not saved)

Key Takeaways

Fine-tuning a saved GPT model restores two pieces of state, the model and the optimizer, both stored in one checkpoint file.
The checkpoint is read with torch.load(), and model_state_dict and optimizer_state_dict feed model.load_state_dict() and optimizer.load_state_dict() respectively.
Both the model and the optimizer must be constructed before their saved states are loaded, because the saved values need an object to fill.
The optimizer is built on the restored model's parameters, with learning rate 5e-4 and weight decay 0.1.
Calling train_simple_function with num_epochs=1 continues training for one more epoch from the restored state.

Related concepts

More concepts starting with T