Open weights

We’ll now see how to load a weights file into our model. For this we'll use the published weights for the GPT-2 model, which is an older model (2019), but it is small and simple enough to work with on a modest laptop.

An important thing to note when loading open weights using your own LLM code is that your LLM needs to be compatible with the architecture and parameter shapes of the LLM that produced the weights. There are many choices you can make when building an LLM: the number of transformer layers, projection dimensions, choice of activation functions, and many more. In our case, we have carefully constructed our LLM to be compatible with GPT-2 weights.

Now the code to actually load the weights into our model looks as follows:

def load_gpt2_small(weights_path, device="cpu"):
    model = LanguageModel(
        vocab_size=50257,
        max_sequence_length=1024,
        model_dim=768,
        head_dim=64,
        num_heads=12,
        hidden_dim=3072,
        num_layers=12,
        qkv_bias=True,
    )

    source_state = torch.load(
        weights_path,
        map_location="cpu",
        weights_only=True,
    )
    translated_state = translate_weights(source_state, model)

    model.load_state_dict(translated_state, strict=True)
    model.eval()
    return model.to(device)

We first construct our model with GPT-2 small's dimensions. The size of every learned tensor depends on these choices, so even one different dimension would make the weights incompatible.

torch.load reads the tensors onto the CPU. The weights_only=True option limits loading to the data types needed for weights, rather than allowing arbitrary Python objects from the file.

Translate the parameter names

A PyTorch weights file is a dictionary from parameter names to tensors. The model.load_state_dict(...) call copies those tensors into the model, matching each dictionary entry to a model parameter by its exact name rather than by its position in the file.

Our teaching model uses descriptive names, while the GPT-2 checkpoint uses shorter ones. Calling load_state_dict with the original names would therefore report missing and unexpected parameters. The tensors already have the right values and shapes; the loader only needs to give them the names expected by our model.

The complete loader applies the same pattern to every parameter in all 12 transformer blocks:

def translate_weights(source_state, model):
    translated_state = {}

    for name, tensor in source_state.items():
        if name.endswith(".att.mask"):
            continue
        translated_state[translate_name(name)] = tensor

    check_keys_and_shapes(translated_state, model.state_dict())
    return translated_state

The checks make sure that no parameter is missing, unexpected, or the wrong shape. Passing strict=True to load_state_dict provides one final safeguard: the program fails clearly instead of silently leaving part of the model randomly initialised.

Generate the next token

Once loaded, the model is used in exactly the same way as before. The difference is that its predictions now come from learned weights rather than random starting values:

model = load_gpt2_small("weights/gpt2-small.pth")
token_ids = torch.tensor(tokenise("The dog fetched the"))

with torch.inference_mode():
    next_token_logits = model(token_ids)
    next_token_id = torch.argmax(next_token_logits).item()

print(detokenise([next_token_id]))

The runnable example keeps all checkpoint-specific work in a separate weight-loading file. This leaves the language-model classes focused on the components we have already explored: embeddings, attention, feed-forward layers, transformer blocks, and the vocabulary projection.