6 — Borrowing GPT-2’s weights

Chapter 5 trained a small GPT on one novel and hit a wall made of data. After a dozen passes over 40,733 tokens the model had learned what it could and started memorizing. Its best loss on unseen text was 5.700. GPT-2 was trained on 40 GB of web text, and doing that again is far beyond one Mac.

It doesn’t have to be done again. In 2019 OpenAI released GPT-2’s trained weights along with its model code, under a modified MIT license. The model built in chapter 4 has GPT-2 124M’s exact shape, so those weights can be poured straight into it. This chapter downloads them, maps every tensor to its place, checks that the result computes the same numbers as an independent implementation, and then looks at what came in: the text it writes, what the weights themselves look like, and what it costs to run.

All runs were on an Apple M3 Pro with 18 GB of memory, with Python 3.12, PyTorch 2.15.0, tiktoken 0.14.0, safetensors 0.8.0, and, for the reference check only, Hugging Face transformers 5.22.2. Seeds are given with each run.

Getting the file

OpenAI’s repository includes a script that downloads the original TensorFlow checkpoint. The same weights are also on Hugging Face as openai-community/gpt2, the account that hosts OpenAI’s original releases there, in the safetensors format. This chapter uses that file:

curl -L -o data/gpt2/model.safetensors \
  https://huggingface.co/openai-community/gpt2/resolve/main/model.safetensors

It is 548,105,171 bytes, with SHA-256 248dfc39…0f81d3a707 at repository revision 607a30d. Safetensors is a deliberately simple format: a short JSON header listing each tensor’s name, data type, shape and byte offsets, followed by the raw numbers. Loading it reads bytes and nothing else. That matters, because the older PyTorch format is built on Python’s pickle, which can run arbitrary code while loading, so a weights file from a stranger is a program from a stranger.

from safetensors.torch import load_file

sd = load_file("data/gpt2/model.safetensors")
print(len(sd))
for name in ["wte.weight", "wpe.weight", "h.0.ln_1.weight", "h.0.attn.c_attn.weight",
             "h.0.attn.c_attn.bias", "h.0.attn.c_proj.weight", "h.0.mlp.c_fc.weight",
             "h.0.mlp.c_proj.weight", "h.0.attn.bias", "ln_f.weight"]:
    print(name, tuple(sd[name].shape))

160
wte.weight (50257, 768)
wpe.weight (1024, 768)
h.0.ln_1.weight (768,)
h.0.attn.c_attn.weight (768, 2304)
h.0.attn.c_attn.bias (2304,)
h.0.attn.c_proj.weight (768, 768)
h.0.mlp.c_fc.weight (768, 3072)
h.0.mlp.c_proj.weight (3072, 768)
h.0.attn.bias (1, 1, 1024, 1024)
ln_f.weight (768,)

load_file returns an ordinary Python dictionary from names to tensors. The names follow OpenAI’s code: wte and wpe are the token and position tables from chapter 2, h.0 to h.11 are the twelve blocks, ln_1 and ln_2 their two layer norms, attn.c_attn the fused query-key-value projection, attn.c_proj the attention output projection, mlp.c_fc and mlp.c_proj the feed-forward network’s two layers, and ln_f the final layer norm. Every block has 13 tensors, and 12 × 13 + 4 = 160.

Two things are notable by their absence and presence. There is no output head: no tensor of shape (50,257, 768) other than wte. That is weight tying from chapter 4; GPT-2’s head is the token embedding. And there is an attn.bias in each block that isn’t a bias at all:

print(sd["h.0.attn.bias"][0, 0, :4, :4])

tensor([[1., 0., 0., 0.],
        [1., 1., 0., 0.],
        [1., 1., 1., 0.],
        [1., 1., 1., 1.]])

It is the causal mask from chapter 3, a 1,024 × 1,024 lower triangle of ones, saved to the file as if it were a weight. Twelve copies of it are 12,582,912 numbers, 50,331,648 bytes or 9.2% of the file, and none of them is learned. The chapter 4 model builds its mask itself and doesn’t save it, so these are skipped. Everything else adds up exactly:

tensors in file:         160
numbers in file:         137,022,720
  of which masks:         12,582,912
  learned parameters:    124,439,808
bytes, parameters only:  497,759,232
bytes, everything:       548,090,880

The parameters number 124,439,808, the same as the tied chapter 4 model, and at 4 bytes each they and the masks account for all but 14,291 bytes of the file, which is the JSON header.

Mapping the names

Most tensors go across one to one. The only real work is a transpose. OpenAI’s code stores each linear layer’s weight as an (input, output) matrix and computes x @ W. PyTorch’s nn.Linear stores (output, input) and computes x @ W.T. So every weight matrix inside a block is transposed on the way in. The embeddings, biases and layer norm parameters are not matrices of that kind and copy as they are.

The fused query-key-value layer from chapter 3 pays off here: GPT-2 stores q, k and v side by side in one 768 × 2,304 matrix, in that order, and chapter 3’s qkv layer splits its output in the same order, so the matrix copies across whole.

The loader builds a tied model and copies each tensor in, refusing any shape mismatch so that a mistake fails loudly instead of producing a quietly broken model:

import torch
from safetensors.torch import load_file

def _copy(dst, src):
    if dst.shape != src.shape:
        raise ValueError(f"shape mismatch: model {tuple(dst.shape)} vs file {tuple(src.shape)}")
    with torch.no_grad():
        dst.copy_(src)

def load_gpt2(path, device="cpu", dropout=0.0):
    sd = load_file(path)
    model = GPT(GPTConfig(dropout=dropout, qkv_bias=True, tie_weights=True))
    _copy(model.tok_emb.weight, sd["wte.weight"])
    _copy(model.pos_emb.weight, sd["wpe.weight"])
    for i, block in enumerate(model.blocks):
        p = f"h.{i}."
        _copy(block.norm1.weight, sd[p + "ln_1.weight"])
        _copy(block.norm1.bias, sd[p + "ln_1.bias"])
        _copy(block.attn.qkv.weight, sd[p + "attn.c_attn.weight"].T)
        _copy(block.attn.qkv.bias, sd[p + "attn.c_attn.bias"])
        _copy(block.attn.proj.weight, sd[p + "attn.c_proj.weight"].T)
        _copy(block.attn.proj.bias, sd[p + "attn.c_proj.bias"])
        _copy(block.norm2.weight, sd[p + "ln_2.weight"])
        _copy(block.norm2.bias, sd[p + "ln_2.bias"])
        _copy(block.ffn.fc.weight, sd[p + "mlp.c_fc.weight"].T)
        _copy(block.ffn.fc.bias, sd[p + "mlp.c_fc.bias"])
        _copy(block.ffn.proj.weight, sd[p + "mlp.c_proj.weight"].T)
        _copy(block.ffn.proj.bias, sd[p + "mlp.c_proj.bias"])
    _copy(model.final_norm.weight, sd["ln_f.weight"])
    _copy(model.final_norm.bias, sd["ln_f.bias"])
    return model.to(device).eval()

.T is the transpose. dst.copy_(src) overwrites a parameter’s numbers in place; the trailing underscore is PyTorch’s convention for an operation that modifies its tensor. Doing it inside torch.no_grad() keeps PyTorch from recording the copy as part of a computation to differentiate. The head never appears: because head.weight is tok_emb.weight, copying wte fills both. Dropout defaults to 0, since the loaded model is for running, not training.

All 148 learned tensors in the file (160 minus the 12 masks) found a place, and every shape agreed on the first try, which is what a model built to the same specification should do.

Checking it computes the same numbers

A shape match only proves the tensors fit. A transposed matrix of the right size, or q and k swapped, would load without complaint and produce a fluent-looking but wrong model. The real test is to run the same input through an independent implementation and compare the outputs.

Hugging Face’s transformers library has its own GPT-2, written separately, and can load the same file:

from transformers import GPT2LMHeadModel
import tiktoken

enc = tiktoken.get_encoding("gpt2")
ours = load_gpt2("data/gpt2/model.safetensors")
ref = GPT2LMHeadModel.from_pretrained("data/gpt2").eval()

idx = torch.tensor([enc.encode("Every chapter runs on one Mac, and every number shown came from a real run.")])
with torch.no_grad():
    a = ours(idx)
    b = ref(idx).logits
print((a - b).abs().max().item(), a.min().item(), a.max().item())
print((a.argmax(-1) == b.argmax(-1)).all().item())

0.0001678466796875 -186.30389404296875 -32.22185516357422
True

.abs().max() is the largest absolute difference anywhere among the 854,369 logits. The logits themselves run from about −186 to −32, so a difference of 0.00017 is a relative error of about one part in a million, the size of rounding differences in 32-bit floating point. Every position’s top prediction agrees.

The two implementations are not bit-identical, and shouldn’t be expected to be: they do the same arithmetic in a different order (the reference uses PyTorch’s fused attention kernel, for one), and floating-point addition gives slightly different results in a different order. A useful tolerance has to sit above that rounding noise and far below any real bug. On a harder input, four sequences of 1,024 random tokens (seed 123), the full context:

The reference ran on the CPU in both rows; only this chapter’s model moved to the GPU, which adds its own slightly different rounding.

A tolerance of 10⁻³ on the logits passes on both devices. A real mapping mistake does not come close. Loading c_attn without its transpose fails the shape check outright, and swapping just the query and key blocks of block 6’s qkv layer, which passes the shape check, changes the logits by up to 46.55 on the sentence above and changes the top prediction at 10 of its 17 positions.

What it writes

With the weights in, chapter 5’s generate works unchanged. Greedy decoding, and then top-k sampling with k = 40 and temperature 1.0 at seeds 1 and 2, 40 new tokens each, on the CPU:

[greedy] The Time Traveller's Guide to the Galaxy, by John W. Campbell, is a collection of essays by the author, and is available for purchase from Amazon.com.

The Time Traveller's Guide to

[top-k 40, seed 1] The Time Traveller: The Mummy" was broadcast on December 20, 2003, about two years before the arrival of Dr Jekyll and Mr Hyde from the moon.

It was also released on DVD in

[top-k 40, seed 2] The Time Traveller Guide to New York City

This week's episode introduces us to Time Travelers, the story of the time people who started their world. We are joined by the BBC's Jonathan Pearce, writer

[greedy] The best way to learn a new programming language is to learn it from a beginner.

The best way to learn a new programming language is to learn it from a beginner. Learn the language from a beginner.

Learn the language from a

[top-k 40, seed 1] The best way to learn a new programming language is through programming. A programmer learns languages by experience and reading through the literature. This can include basic concepts like inheritance, inheritance analysis, and inheritance semantics. A programmer who is able to program with those concepts

[top-k 40, seed 2] The best way to learn a new programming language is to learn it at school. You might be surprised at how much it influences your life.

The best way to learn Python is to learn it at school.

Pilots can learn

[greedy] In 1969, the first of the two major U.S. military exercises, the "Operation Desert Storm," was held in the Mojave Desert. The first of the two major U.S. military exercises, the "

[top-k 40, seed 2] In 1969, the first real attempt to build that was attempted in the early 1960s when John Van Pelt joined the club from the Australian Football League. He then joined the club in 1975 before getting into the role of manager

The jump from chapter 5 is enormous. These are grammatical sentences that hold a topic over several clauses, with plausible names, dates and places. But look at what they are. Asked about the Time Traveller, it invents a guide book attributed to a real science-fiction editor and says it is for sale on Amazon. Asked about 1969, the greedy version puts Operation Desert Storm, which happened in 1991, in the Mojave Desert. Asked for advice about programming languages, it repeats itself. GPT-2 continues text in the style of the web pages it was trained on. It was never trained to be correct or to answer, only to be likely.

On held-out text

The fairest comparison with chapter 5 is the same measurement: the loss on the held-out last 10% of The Time Machine, in windows of 128 tokens.

GPT-2 never saw this split of the book during anything in this track (whether some copy of the novel was in WebText is unknown), yet it predicts the held-out chapter far better than a model that trained on the rest of the book until it memorized it. Given 1,024 tokens of context instead of 128, it does better still, because more of the earlier text is available to every prediction. (The held-out text holds only 4 full windows of 1,024 tokens, so that figure rests on 4,096 predictions rather than all 4,480 of the 128-token windows, and is the less precise of the two.) Its loss on the training split, 4.263, is no lower than on the validation split, so there is no sign it has the book memorized.

What the numbers in the weights look like

The 124 million numbers are not mysterious individually. Across all of them, the mean is −0.0005 and the standard deviation 0.135; 59% of them are smaller than 0.1 in magnitude. A typical weight matrix looks like a bell curve centered on zero:

The averages hide the interesting parts, which are in the extremes. The largest magnitude anywhere is 17.42, in the scale of the final layer norm, which ranges from 0.004 to 17.42: after twelve blocks, a handful of the 768 dimensions are amplified far more than the rest. The residual stream itself shows the same thing. Running the first 256 tokens of The Time Machine through the model and averaging the size of each of the 768 numbers in the residual stream after block 7 (leaving out the first position), the median dimension averages 1.74 in magnitude, but dimension 447 averages 34.1, dimension 64 27.4 and dimension 373 17.9. A few fixed directions carry values twenty times larger than the rest. Large-magnitude dimensions like these are a known feature of trained transformers, and they matter in practice: they are one reason squeezing a model into fewer bits per weight is harder than it looks.

The token embedding table is where chapter 2 first looked at these weights, finding " seven" next to " six" and " eight". A few more tokens, with their closest rows by cosine similarity:

Nothing told the model that " Paris" and " France" are related, or that " king" and " queen" are. The rows ended up near each other because those tokens are used in similar contexts, and predicting the next token rewarded the model for treating them alike. The case and leading-space variants of each word, which chapter 1 showed are separate tokens, land close together too.

The position table shows its own pattern. Comparing each position’s vector with every other position’s, every 32nd position:

Neighboring positions have nearly identical vectors (position 100 and 101 have a cosine similarity of 0.999), and similarity falls off and then oscillates with distance, in bands, like a set of waves of different lengths. The original transformer used fixed sine and cosine waves for positions. GPT-2 learned its position vectors from scratch and ended up with something wave-like anyway. Position 0 stands apart from the rest, with a vector more than twice as long (norm 9.88 against 4.06 for position 5). The first position is special in a way no other is: a token there has nothing to attend to but itself, and a plausible reading is that the model learned a distinctive marker for it.

Size and speed on this Mac

In 32-bit floats the model is 497.8 MB of parameters. load_gpt2 takes 2.49 seconds. Of that, load_file takes only 0.01 seconds, because safetensors maps the file into memory instead of reading and decoding it up front; the rest goes into building a randomly initialized model and copying the real weights over it.

Timed with no other model running on the machine, after two warm-up passes, median of five for the forward pass and a single run for generation:

A forward pass over a full 1,024-token context takes 0.155 seconds on the GPU. Generation pays for a whole forward pass per new token, about 20 milliseconds each on the GPU here, because chapter 4’s generate runs the entire sequence through the model again every time; chapter 9 fixes that with a cache. The model also runs in 16-bit floats on MPS, at half the memory, with 272 MB allocated. It is not free: on the same four random sequences, the 16-bit logits differ from the reference by up to 0.294, and the top prediction changes at 2.8% of positions. Random tokens are a harsh test, since they leave many near-ties between candidates, but it shows that lower precision is a trade, not a free saving.

Where this leaves things

The model from chapter 4, filled with OpenAI’s weights, is GPT-2: its logits match Hugging Face’s implementation to within 0.0003, it predicts unseen text far better than anything trainable here from scratch, and it runs comfortably on a laptop. Its embeddings put Paris near France and king near queen.

But it only does one thing: continue text the way text on the web tends to continue. Asked a question, it writes something shaped like what follows such a question on a web page, which might be an answer, a second question, or an advertisement. It invents facts with confidence. It loops. It is fluent and general, and not useful for any specific job.

The cheapest way to make it useful is to point it at one narrow task where the right answer is clear and can be checked. Chapter 7 takes this model, replaces its output head with one that has just two outputs, and trains it to tell spam text messages from real ones.