Saving a PyTorch Model

After fine-tuning, a model's learned values exist only while the program is running. Once the program ends, those values are gone unless they were written to a file. Saving is the final step of the fine-tuning workflow: the fine-tuned model is written to disk so it can be reused in future projects, and loaded back into a program later.

This article follows one save and one load from start to finish. As you read, trace each line and notice which value it produces before the next line uses it. Most save and load errors come from a step that used a value the previous step did not produce in the form expected.

The saving work happens in one call: torch.save(model.state_dict(), file_name). The first argument, model.state_dict(), gives the model's state dictionary. The second argument, file_name, is a string that names the file on disk. The model is saved as a .pth file. The file name is computed on the line before the save, so the name is already fixed before torch.save runs.

The diagram shows that the state dictionary is taken from the model, handed to torch.save, and written to the file named by file_name. The model object itself is not what gets written; the save operates on the state dictionary. The file name is the only thing that determines where the output goes, so a wrong file name produces a file with the wrong name rather than a missing save.

The file name is not typed by hand. It is generated from the chosen model name. The code takes CHOOSE_MODEL, removes every space, opening parenthesis, and closing parenthesis using re.sub with the pattern [ ()], and then appends the suffix -sft.pth. The result for the course model is gpt2-medium355M-sft.pth.

A state dictionary is a mapping. Each entry pairs a name with a value, so you can look up a stored value by its name. When torch.save writes the state dictionary, the names and the values are both written to the file. When the file is read back, load_state_dict uses those names to decide which saved value belongs to which place in the model.

The example keys weight and bias are illustrative. The important point is structural: every name in the saved mapping must correspond to a place in the model that receives a value. If the names in the file do not match the names expected by the model, the saved values cannot be placed correctly. This is why the same model structure has to be used on both sides of the save and load.

The call saves model.state_dict(), not the model object. These are different things, and the difference determines what the .pth file contains and what you must have ready when you load it.

Loading runs in the opposite direction to saving. The file is read by torch.load, the resulting state dictionary is passed to load_state_dict, and the model receives the saved values. The saved model can be loaded later with the following call, which uses the same file name that the save produced.

Read the sequence in order. The file produces nothing until torch.load reads it. Its output is a state dictionary, and that dictionary is what load_state_dict receives. Only at the final step does the model's own storage change. If you reverse any two of these steps, the program is describing a different operation than the one the source intends.

When a save or load fails, trace the file name first. Both steps must use the same generated string, so a mismatch is the most common cause of a broken load. Work through the checks below in order.