Inspecting the Pretrained GPT-2 Parameter Dictionary
assign compares the shapes of the destination and proposed value, raises a ValueError on a mismatch, and otherwise returns the right-hand value as a trainable PyTorch parameter.
Why a Fresh GPTModel Needs Pretrained Values
A GPTModel starts out with randomly initialized weights. On its own, that model has not learned anything useful. To make it behave like the pretrained GPT-2 model, its weights must be replaced with values taken from a pretrained source. In this lesson, that source is a Python dictionary called params. Each entry in params holds a set of values, and the loading step copies those values into the matching places in the model.
Loading weights is easy to get wrong because the loader makes many assignments, and any one of them can pair the wrong pretrained value with the wrong model parameter. Before following the assignments, it helps to ask a debugging question: if a value from params does not fit the model parameter it is being assigned to, where should the problem become visible? Keep that question in mind as you read. The answer is the central idea of this lesson.
What do you think happens?
If a params value has a different shape from the model parameter it is being assigned to, where should the problem be caught?
Reveal answer
At the assign call for that pair, which compares the two shapes and raises a ValueError on a mismatch.
The loader does not silently copy a value of the wrong size. Each assignment passes through assign, which stops the load right at the incompatible pair.
Tracing Weights from params into the Model
Think of load_weights_into_gpt as a set of deliberate correspondences. The function takes a GPTModel instance and the params dictionary. It then assigns pretrained values to matching locations in the model. It handles the positional and token embeddings, visits each transformer block, and assigns the final normalization parameters and the output-head weight. Nothing is copied by guessing. Each destination in the model is paired with a specific entry in params.
Read the diagram from top to bottom. The params dictionary is the source. Each branch selects a group of entries that belong together: token embeddings, combined attention entries, transformer block entries, and final normalization entries. Each entry then lands in a specific GPTModel location. The token embedding entry lands in two terminals, which is the case covered later in this lesson.
What assign Does With Each Pair
The loading function makes many assignments, so it relies on a small guard against incompatible pairs. The guard is assign. It receives the model's current parameter as left and the proposed pretrained value as right. Before anything else, it compares their shapes.
assign compares the shapes of the destination (left) and the proposed value (right). If the shapes differ, it raises a ValueError and the load stops. If the shapes match, it returns the right-hand value as a trainable PyTorch parameter.
Tracing one assignment through assign
The loader pairs a model's current parameter with a pretrained value. Follow the call and decide whether it returns a parameter or stops with an error.
- 1Identify the pair
left is the model's current parameter. right is the proposed pretrained value taken from params.
- 2Compare shapes
assign checks whether left and right have the same shape before doing anything else.
- 3Take the branch
If the shapes differ, assign raises a
ValueErrorand the load stops at this pair. If they match, assign continues. - 4Prepare the parameter
On a match, assign returns the right-hand value as a trainable PyTorch parameter, ready to take the place of the model's weight.
A matching pair produces a trainable parameter. A mismatched pair stops the load at that exact assignment, with the error naming both shapes.
Weight Tying Between Embedding and Output Head
The token embedding entry, params['wte'], is assigned in two places. It goes to gpt.tok_emb.weight, which is the input embedding, and it also goes to gpt.out_head.weight, which is the output head. Using one set of values for both places is called weight tying. The source identifies this as a feature of the original GPT-2 model, and notes that it reduces the total number of parameters.
The diagram shows the single source entry feeding two destinations. This is the reason the loader makes two assignments from the same params key. Both destinations receive the same pretrained values, so both locations are filled from one place in the dictionary.
Finding a Shape Mismatch
When loading fails, start with the assignment that raised the shape-mismatch error. The error reports the left and right shapes, which gives you a concrete place to inspect the proposed pairing. The shape check is useful because it stops the load at the incompatible pair instead of letting a wrong value slip into the model unnoticed.
A shape check has a limit. It confirms that two values fit together structurally. It cannot confirm that the selected values have the correct meaning. If the shapes match but the resulting model cannot produce coherent text, the mapping may still be wrong. In that case, revisit the source-to-destination correspondence rather than assuming the check proved it correct.
Common mistake Assuming matching shapes prove the weights are correct
Why it failsassign only compares shapes. Two tensors can have the same shape and still come from the wrong params entry.
Do this insteadAfter a successful load, check that each params entry was paired with the intended GPTModel component.
Common mistake Matching params names to model names by guess alone
Why it failsThe source notes that naming conventions differ between the two implementations, so guessed matches can be unreliable.
Do this insteadConfirm each source-to-destination correspondence explicitly, following the assignments in load_weights_into_gpt.
Common mistake Forgetting the token embedding feeds the output head
Why it failsIf you only track the input embedding, you will miss that params['wte'] also fills gpt.out_head.weight.
Do this insteadTreat params['wte'] as a two-destination entry when you trace the load.
Common mistake Skipping the transposes and splits
Why it failsAttention weights are split into query, key, and value parts, and some arrays are transposed before assign. Skipping either step changes what is being paired.
Do this insteadTrace the split and the transpose as part of the correspondence, not as optional details.
Practicing the Load Trace
A colleague says the pretrained model loads cleanly because no error was raised. Using the trace you have built, list the checks that were actually performed and the checks that were not. Then describe one thing you would still verify before trusting the model's output.
Hints
- Think about what assign compares and what it does not compare.
- Recall the two destinations that params['wte'] fills.
- Consider what happens when the shapes match but the values are paired incorrectly.
Key Takeaways
- A GPTModel starts with random weights, and
load_weights_into_gptreplaces them with values from the params dictionary by assigning each pretrained entry to a matching model location. - assign compares the shapes of the model's current parameter and the proposed value, raises a
ValueErroron a mismatch, and otherwise returns the value as a trainable PyTorch parameter. - The params['wte'] entry is assigned to both gpt.tok_emb.weight and gpt.out_head.weight, which is weight tying and reduces the total number of parameters.
- A shape mismatch is caught at the specific assign call for that pair, which makes the error easy to locate, but a matching shape does not prove the mapping is semantically correct.