Fine-Tuning Pretrained Models
The GPT-3 authors did not share their training dataset, so Dolma is offered as a comparable publicly available resource for pretraining research.
When the Training Data Is Missing
Suppose you want to understand how a large language model was built, but the people who built it have not released the material it learned from. That is the situation with GPT-3. The authors of the GPT-3 paper did not share its training dataset.
This creates a practical problem for anyone who wants to study, reproduce, or research models at this scale. The data that shaped the model is not available to inspect. A useful first question is what a researcher can reach for when the original data is missing, and the second is what costs come with building a model of this kind at all. The sections that follow take those questions in order.
Dolma as a Comparable Public Resource
When the original data is unavailable, a publicly available resource can stand in for it. The one identified here is Dolma, described in the paper 'Dolma: An Open Corpus of Three Trillion Tokens for LLM Pretraining Research' by Soldaini and colleagues, published in 2024. Dolma is designed for research into pretraining large language models, and it contains three trillion tokens.
The word comparable carries the meaning here. Dolma is not presented as the dataset GPT-3 was trained on. It is a public resource that offers a similar kind of material and scale, so that researchers can work with pretraining data without needing access to the private collection.
Why this matters: a stand-in lets a researcher practice the pretraining side of the workflow, inspecting and working with data of a similar kind and scale, even when the exact training material is locked away. The stand-in does not reproduce GPT-3's data, and it should not be described as if it did.
Usage Limits on Public Data
Being publicly available does not mean a dataset can be used without thought. The collection may contain copyrighted works. The exact usage terms may also depend on the intended use case and on the country involved.
The Cost of Pretraining
Pretraining large language models requires access to significant resources and is very expensive. A concrete reference point is the GPT-3 pretraining cost, estimated at $4.6 million in cloud computing credits.
Two details in that sentence deserve attention. The figure is an estimate, not a fixed invoice. It is also expressed in cloud computing credits, which tells you the cost was framed as the compute resources consumed to train the model.
| Stage | What it produces | Cost information in the source |
|---|---|---|
| Pretraining | A pretrained base or foundation model | Very expensive; GPT-3 estimated at $4.6 million in cloud computing credits |
| Fine-tuning | A model adapted to a downstream task, starting from a saved point | No cost figure is given in the source |
Pretraining and fine-tuning compared. The source gives a cost estimate only for pretraining.
Because pretraining is so expensive, a team has a strong reason to ask whether it needs to pretrain at all. The source supports one part of that reasoning directly: a pretrained model can be fine-tuned on downstream tasks, and that is why such models are called base or foundation models.
Base Models and the Fine-Tuning Step
A pretrained large language model is a model that has first been trained on a large dataset. It is also called a base model or a foundation model, names that point to its role. Because it has already been pretrained, it is versatile enough to be further fine-tuned on downstream tasks, meaning the specific jobs a team actually wants done.
Read the pipeline from left to right. The expensive step is pretraining, which produces the base model. Fine-tuning comes later and connects the base model to a particular downstream task. The pretrained model is the common starting point, and the specialized behavior is added afterward.
Tracing the Saved Checkpoint
Training a GPT model is expensive, so you rarely want to throw away work already done. Fine-tuning a pretrained model means picking up from a saved point and continuing training for more passes over the data. The goal is to see exactly what has to be restored so that continued training starts from the saved point rather than from scratch.
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 is torch.optim.AdamW in this workflow. The optimizer keeps its own internal state alongside the model. If you restored only the model, the optimizer would behave as though it had never run.
The loading code reads the checkpoint into a single object, then pulls two entries out of it by key. The key model_state_dict belongs to the model, and the key optimizer_state_dict belongs to the optimizer. These two keys form the whole contract between saving and loading, so they are the first place to look when something goes wrong.
import torch
checkpoint = torch.load(checkpoint_path)
# model and optimizer are rebuilt first, then loaded (see the next section)
model = GPTModel(GPT_CONFIG_124M)
model.load_state_dict(checkpoint["model_state_dict"])
optimizer = torch.optim.AdamW(model.parameters(), lr=5e-4, weight_decay=0.1)
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
train_simple_function(..., num_epochs=1)- 1Line 3 reads the checkpoint file. Nothing is restored yet; the file is just loaded into one object.
- 2Lines 6 and 7 create a fresh GPTModel from GPT_CONFIG_124M, then load the
model_state_dictentry into it. - 3Lines 9 and 10 create the AdamW optimizer on the parameters of the freshly loaded model, then load the
optimizer_state_dictentry into it. - 4Line 12 continues training for one more epoch from the restored state.
torch.load()reads the checkpoint into a single dictionary-like object- The
model_state_dictentry goes tomodel.load_state_dict() - The
optimizer_state_dictentry goes tooptimizer.load_state_dict() train_simple_functionis called withnum_epochs=1; its other arguments come from the earlier chapter
Predicting the Checkpoint Wiring
What do you think happens?
The loaded checkpoint is a dictionary with two entries you care about. Which entry should go to the model, and which should go to the optimizer?
Reveal answer
The model_state_dict entry goes to model.load_state_dict(), and the optimizer_state_dict entry goes to optimizer.load_state_dict().
Each entry has one destination, and the keys are the contract between saving and loading. One call cannot serve both because the model and the optimizer hold different state: the model's learned parameters and the optimizer's own internal state.
Rebuilding Before Restoring
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 follows the same pattern: it is constructed first, then given its saved state. 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. The order of construction matters, because each loading call needs an existing object to load into.
Running One More Epoch
Once the model and optimizer are restored, the final step is to continue training. Call train_simple_function with num_epochs=1. That single call trains the model for another epoch, starting from the restored state rather than from random initial values.
The internals of train_simple_function are not described here. This article only shows where the call sits in the resume sequence: after the checkpoint has been read and both objects have been rebuilt and restored.
Resuming Training From a Saved Checkpoint
A checkpoint file contains the keys model_state_dict and optimizer_state_dict. Restore the GPT model and its optimizer, then continue training for one more epoch.
- 1Read the checkpoint
Call
torch.load()on the checkpoint file to get one object holding both saved entries. - 2Rebuild the model
Create a fresh GPTModel from GPT_CONFIG_124M so there is a structure to receive the saved parameters.
- 3Restore the model
Pass the
model_state_dictentry tomodel.load_state_dict(). - 4Rebuild the optimizer
Create torch.optim.AdamW on the parameters of the restored model, with learning rate 5e-4 and weight decay 0.1.
- 5Restore the optimizer
Pass the
optimizer_state_dictentry tooptimizer.load_state_dict(). - 6Continue training
Call
train_simple_functionwithnum_epochs=1, which trains one more epoch from the restored state.
The model and optimizer both resume from the saved point, and one additional epoch of training is run.
Mistakes When Resuming Training
Common mistake Restoring only the model and skipping the optimizer
Why it failsThe optimizer keeps its own internal state alongside the model. Without restoring it, the optimizer behaves as though it had never run.
Do this insteadRestore both model_state_dict and optimizer_state_dict from the same checkpoint.
Common mistake Loading a saved state before constructing the object that receives it
Why it failsThe saved state is only a set of values. A load call fills an existing model or optimizer, so the object must be created first.
Do this insteadBuild GPTModel from GPT_CONFIG_124M and the AdamW optimizer on the model's parameters, then call the load methods.
Common mistake Passing the wrong entry to the wrong loading call
What it looks likePassing optimizer_state_dict to model.load_state_dict()Why it failsEach entry belongs to one destination. The keys are the contract between saving and loading.
Do this insteadSend model_state_dict to model.load_state_dict() and optimizer_state_dict to optimizer.load_state_dict().Common mistake Treating a public dataset label as blanket permission to use it for any purpose
Practice Resuming a Checkpoint
Without running any code, trace a resume script. Name the checkpoint key that feeds model.load_state_dict(), the key that feeds optimizer.load_state_dict(), and the step that must come before each load. Then state where a key mismatch would be reported.
Hints
- Look for the two keys the checkpoint dictionary is read by.
- Think about which object must exist before a load call can fill it.
- A mismatch would surface at the load call that receives the entry with the wrong key or shape.
Key Takeaways
- A missing training dataset can be partly addressed by a comparable public resource such as Dolma, which is not the same as the original data, and public datasets may still carry usage limits that depend on purpose and country.
- Pretraining is very expensive, with the GPT-3 estimate at $4.6 million in cloud computing credits; a pretrained base or foundation model can then be fine-tuned for downstream tasks.
- Resuming training requires restoring two saved components from one checkpoint: the model, through
model_state_dict, and the AdamW optimizer, throughoptimizer_state_dict. - The model and optimizer are constructed before their saved states are loaded, because a load call fills an existing object.
- Calling
train_simple_functionwithnum_epochs=1 continues training for one more epoch from the restored state.