Saving and Loading Model Checkpoints

Training a large language model is a long process, and a training run does not always finish in one sitting. Training a GPT model is expensive, so you rarely want to throw away the work already done. Instead of treating every training session as a fresh start, the progress made so far can be stored and then picked up again later. This picking up is called fine-tuning a saved model: you continue training for more passes over the data, starting from a 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. This article unpacks what that sentence requires, one piece at a time, and explains why the order of the steps matters. The central question 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?

Saving is the act of writing the current weights to storage. Saving model weights stores the learned values of a model so that they can be used later. Loading is the reverse: reading previously saved weights back in, so the model starts from that stored state rather than from scratch. The source gives both operations the same purpose, which is to continue training an LLM.

The two operations work best when you think of them as a pair. A save without a later load stores the progress but never puts it back to work. A load is only meaningful when there is a saved set of weights to restore. The diagram below traces the weights through one full break in training, from their random starting values to one extra epoch after they have been restored.

The important comparison is between the restored weights and a randomly initialized model. A randomly initialized model starts with arbitrary values and has learned nothing yet. The restored model carries the values that were learned before the break. When training resumes from the restored state, it continues from the point where the saved weights left off rather than from those random initial values.

A saved checkpoint holds two separate pieces of state, both stored together in one checkpoint file. 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. The optimizer keeps its own internal state alongside the model. Restoring only the model would leave the optimizer behaving as though it had never run, so the optimizer's saved state has to come back too.

The loading code reads the checkpoint into a single object using torch.load(). That object is a dictionary, and you pull 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 are the whole contract between saving and loading. If something goes wrong when you resume, they are the first place to look.

Each entry must reach a specific destination, and one call cannot serve both. The entry model_state_dict feeds model.load_state_dict(), which fills the model. The entry optimizer_state_dict feeds optimizer.load_state_dict(), which fills the optimizer. The mapping is one entry to one call, and it is the wiring you need to hold in your head before reading any resume code.

A mismatch shows up at the loading call that receives the wrong entry. If a key is wrong, or the entry does not match the object it is loaded into, the failure appears where that load call runs, not earlier in the script. That is why the two keys and the two load calls are the places to check first.

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, 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 matters because the optimizer is built from the model's parameters, so the model must already exist at that point.

Once the model and optimizer are restored, the final step is to continue training. The source calls for train_simple_function with num_epochs=1. That single call trains the model for another epoch, starting from the restored state. Over that one epoch, the weights are updated from their starting values, which are the restored values, toward whatever values the continued training produces at its end.