Saving a GPT Model and Optimizer Checkpoint

Training a language model produces state that is worth keeping. The model's learned values change as training runs, and the optimizer's bookkeeping changes with them. A practical workflow therefore saves both to a file and reads them back later, instead of retraining from nothing each time.

This article starts from that saved file. The task is to rebuild a GPT model and its optimizer from a checkpoint, and then to compute loss values on the training data and on the validation data. The approach is code first: you read the loading steps as a trace, predict what each step produces, and then check your prediction against the explanation.

The loading procedure uses five statements, and each one either reads from the checkpoint, builds an object, or attaches saved values to an object. Knowing which of the three a line does is the key to finding where a run goes wrong. Here is the sequence in order.

The source describes the loading code as a flow of values. A file on disk is read into a Python object. That object holds saved entries under named keys. Two of those entries are then used to fill in two live objects, the model and the optimizer. The checkpoint is not the model, and the model is not the optimizer. The file is a container of named entries; the live objects are what the program builds and then fills.

The saved object holds two entries that matter for this procedure. The entry under model_state_dict feeds the GPT model, and the entry under optimizer_state_dict feeds the AdamW optimizer. Each key selects one saved group of values, and each group goes into exactly one live object.

Read the diagram from top to bottom. The file is the single source. Two keys pull two different groups out of it. Each group lands in a different live object. Nothing in the diagram crosses over: the model entries never go to the optimizer, and the optimizer entries never go to the model.

Loading is a two-part act. First, the program constructs a live object. Second, a saved entry is poured into that object. The order matters. The source calls load_state_dict on an object that already exists, so the object has to be built before any saved state can be attached to it.

Notice that the optimizer is built over model.parameters(), so the model must already exist before the optimizer can be constructed from it. The optimizer's saved state is then loaded last, after the optimizer object is in place.

Once the model is restored, the source calculates two loss values for the GPT model: one for the training set and one for the validation set. Both are computed with a custom function named calc_loss_loader. The function is called twice, with the same model and the same device. Only the data loader changes between the two calls.

Because the model and device are held fixed, any difference between train_loss and val_loss comes from the data each call was given, not from a different model or a different device. The two values are separate results of the same function applied to two different inputs.

A run can fail in the loading path or in the loss-calculation path. The source points to specific checks that narrow the search. Work through the lines in order and ask, for each one, whether it reads, builds, or attaches.