Loading Pretrained Weights

A published checkpoint is a file of named tensors, and the names, shapes, and layouts come from the implementation that wrote it, not from yours. The pretrained file has h.0.attn.c_attn.weight of shape (768, 2304). Your model expects blocks.0.attn.qkv.weight of shape (2304, 768), because nn.Linear stores its weight as (out, in) and the published file stores the transpose. One parameter carries three identities (name, shape, byte layout), and the loader has to reconcile all three.

A loader that copies blindly puts the right tensor in the wrong place, and you get a model that generates nonsense. A loader that refuses to copy on a shape difference but logs nothing leaves you guessing which tensor failed to land. The loader here is explicit: every assignment is logged, every shape is checked, and a LoadReport summarizes hits, misses, and shape mismatches so you can read what happened.

Loading pretrained weights means reading each named tensor from a file, translating its name to a local parameter name, checking its shape, and copying it into the model. The name mapper is just a function from string to string. The shape check is one if. The assignment happens inside torch.no_grad() so autograd does not track the load, and the report holds the outcome of every name.

The GPT-2 naming convention

Published GPT-2 weights live under names like:

Two things to plan for. First, the c_attn, c_proj, and c_fc linears are stored transposed relative to what nn.Linear.weight expects, so the loader transposes during assignment. Second, the LM head is not in the file at all. The GPTModel constructor already ties the head to tok_embed, so loading wte fills both, and the loader re-aliases after the load only if the two data pointers differ.

The local naming convention

The model in this track uses descriptive names:

The mapping is a fixed function, shipped as a dict that the loader consults for each name.

The stub fixture

Real GPT-2 weights are 0.5 GB, so the demo does not download them. It generates a small safetensors fixture on first run, using the exact GPT-2 naming convention and shapes for a 4-block model (d_model 192, 6 heads, vocab 256, context 64) instead of the 12-block, 768-wide original. The fixture has the right structure to exercise every code path in the loader. Swap in the real file and the loader works without modification.

code/main.py implements:

• A small replica of the lesson 35 GPTModel, so this lesson stands on its own.
• make_pretrained_to_local(num_layers), which expands the per-layer entries.
• load_safetensors(model, path), which iterates names, maps them, checks shape, transposes the conv1d-style weights, and assigns under torch.no_grad(). It returns a LoadReport.
• make_stub_safetensors(path, cfg), which generates a fixture file with the exact pretrained naming convention.
• A demo that creates outputs/gpt2-stub.safetensors on first run, builds a fresh model, captures one generated continuation from random init, loads the stub, captures another continuation, prints both, and verifies the two differ (the load actually changed the model).

Run it:

python3 code/main.py

Output: the fixture path, a per-name load log, a LoadReport summary, a continuation before the load, a continuation after the load, and and a shape mismatch check: the demo writes a separate outputs/gpt2-bad.safetensors whose wte.weight has the wrong width, loads it into a fresh model, prints the report, and deletes the file. The report shows the mismatch, nothing is assigned, and the model's parameters are listed under missing. Only a missing file raises (FileNotFoundError).

• safetensors for the on disk format and a streaming reader.
• torch for the model and the assignment math.
• No transformers, no huggingface_hub, no network calls.

Three patterns let a loader survive weights you did not create.

Validate the file before any assignment. Open the file, list every tensor name with its dtype and shape, run the full mapping with shape checks, and start assigning only on success. A half-loaded model fails silently.

Log each assignment with the source name and the destination name. When something looks wrong, the log tells you which tensor landed where, and the alternative is reading hexdumps. The LoadReport dataclass here tracks loaded, missing, unexpected, and shape_mismatch lists and prints a summary at the end.

Treat the LM head as a weight tying alias, not a separate copy. The constructor sets self.lm_head.weight = self.tok_embed.weight, and the loader re-aliases after the load only if the data pointers differ. Copying the embedding matrix into a fresh lm_head.weight parameter breaks tying and quietly doubles your parameter count.

• The loader works for any safetensors file that uses the pretrained naming convention. Real GPT-2 files (small / medium / large / xl) load without code changes; only the model config differs.
• The same pattern extends to LLaMA, Mistral, and Qwen weights once you update the name map. The shape checks and the report stay identical.
• Sanity generation after a load is a quick gate. If the post-load samples look like the pre-load samples, the load did not change the model, which means the mapping silently missed every tensor.

1. Add a dtype argument to the loader that casts each tensor to a target dtype (bfloat16, float16, float32) during assignment. Confirm a float32 model can be downcast to bfloat16 and still generate.
2. Add an expected_layers argument that refuses to load a checkpoint whose h.N indices do not match the model's num_layers.
3. Plug the loader into the lesson 35 generation function and produce two side by side samples: one from random init, one from the loaded fixture.
4. Add an export path: write the current model state into a fresh safetensors file using the pretrained naming convention. Round trip the loader and confirm the report has zero shape mismatches.
5. Extend make_pretrained_to_local to handle the LLaMA naming convention (no biases, RMSNorm, fused qkv layout) and re-run the loader on a stub LLaMA fixture you generate.

• Phase 19 lesson 35 for the architecture that receives the weights.
• Phase 19 lesson 36 for the training loop that produces a checkpoint of the same shape.
• Phase 10 lesson 11 (quantization) for what to do with the loaded weights when memory is tight.
• Phase 10 lesson 13 (building a complete LLM pipeline) for the full lifecycle around load and inference.