Loading the Weights

Every check so far ran on tiny random weights. That was the point — correctness is a property of the arithmetic, and random numbers exercise it in a millisecond. But the engine in Chapter 1 did not print random text; it continued a sentence about capitals, because it ran on the real trained Qwen3-0.6B weights. A model is two things: the architecture, which you have now built line by line, and the trained parameters, which you have not. This chapter is how the second gets into the first.

A released checkpoint is not a program. It is a dictionary — names to tensors, model.embed_tokens.weight mapped to a matrix, and hundreds more. Loading is the careful business of matching each of those names to the module in your model that should hold it. The danger is that the match is easy to get wrong and easy to not notice: send a tensor to no module, or the wrong one, and the layer keeps the random values it was born with, and the model runs anyway — producing fluent, confident nonsense with no error to point at.

9.1 Two names for the same tensor

The checkpoint you download names its tensors the way HuggingFace’s Transformers library does — model.layers.0.self_attn.q_proj.weight, mlp.gate_proj.weight, input_layernorm.weight — because that is the format the weights are released in. The modules you built name them forge’s way — trf_blocks.0.att.W_query.weight, ff.fc1.weight, norm1.scale. Same trained numbers, different labels. Something has to translate, and that translation is a fixed table.

HF_BLOCK_MAP = {
    "self_attn.q_proj.weight": "att.W_query.weight",
    "self_attn.k_proj.weight": "att.W_key.weight",
    "self_attn.v_proj.weight": "att.W_value.weight",
    "self_attn.o_proj.weight": "att.out_proj.weight",
    "self_attn.q_norm.weight": "att.q_norm.scale",
    "self_attn.k_norm.weight": "att.k_norm.scale",
    "mlp.gate_proj.weight": "ff.fc1.weight",
    "mlp.up_proj.weight": "ff.fc2.weight",
    "mlp.down_proj.weight": "ff.fc3.weight",
    "input_layernorm.weight": "norm1.scale",
    "post_attention_layernorm.weight": "norm2.scale",
}


def remap_from_hf(hf: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
    out: dict[str, torch.Tensor] = {
        "tok_emb.weight": hf["model.embed_tokens.weight"],
        "final_norm.scale": hf["model.norm.weight"],
        "out_head.weight": hf.get("lm_head.weight", hf["model.embed_tokens.weight"]),
    }
    layer = 0
    while f"model.layers.{layer}.self_attn.q_proj.weight" in hf:
        for src, dst in HF_BLOCK_MAP.items():
            out[f"trf_blocks.{layer}.{dst}"] = hf[f"model.layers.{layer}.{src}"]
        layer += 1
    return out

Listing 9.1 · the name table and the function that walks it. HF_BLOCK_MAP pairs each per-block tensor’s HuggingFace name with its forge name; remap_from_hf applies it to every layer and rebuilds the whole state dictionary under forge’s names.

Read remap_from_hf top to bottom. Three tensors sit outside any block and are mapped by hand: the token embedding, the final norm, and the output head. That last line carries a real decision — hf.get("lm_head.weight", hf["model.embed_tokens.weight"]). Some checkpoints ship a separate output projection; many, including this one, tie it to the embedding — the same matrix reads tokens in and scores them out — and simply omit lm_head.weight. The .get with a fallback handles both: use the released output head if it exists, otherwise reuse the embedding. Then the loop walks the layers, and for each one applies the eleven-entry block map, until it runs out of layers. What comes back is a dictionary under forge’s names, ready to load.

9.2 The load, and the guard against silence

With the names translated, the tensors go into the model. That is one PyTorch call — but the way forge makes it is the whole point of this section.

def load_checkpoint(model: nn.Module, path: Path) -> None:
    state = cast("dict[str, torch.Tensor]", torch.load(path, map_location="cpu", weights_only=True))
    missing, unexpected = cast(
        "tuple[list[str], list[str]]", tuple(model.load_state_dict(state, strict=False))
    )
    missing_parameters = [k for k in missing if "cos" not in k and "sin" not in k]
    if missing_parameters or unexpected:
        raise ValueError(
            f"state dict mismatch: missing {missing_parameters}, unexpected {unexpected}"
        )

Listing 9.2 · load the file, copy the tensors in, then refuse to continue if anything did not match.

torch.load reads the file into a state dictionary. load_state_dict(..., strict=False) copies each tensor whose name matches a module in the model, and — because strict is off — returns two lists instead of raising: the model parameters that no tensor filled, and the checkpoint names that matched no module. That return is the whole point. A strict load (strict=True) raises PyTorch’s own generic error the instant either list is non-empty and hands you nothing to examine. forge wants the lists themselves, so it can judge them and raise an error that names exactly what went wrong.

And it judges them strictly. The one subtlety is the cos and sin tables from Chapter 4: those are non-persistent buffers — computed at construction, and by definition not part of the state dictionary at all, so they are never written to a checkpoint and never turn up as missing. forge filters them out of the missing list anyway, a belt-and-braces guard for the day someone makes them persistent; today that filter removes nothing. What carries the weight is the line after it: if either list still holds a name, forge raises. A misnamed key, a checkpoint for the wrong architecture, a dropped tensor — each strands a name in one list or the other, and each becomes a loud ValueError naming exactly what did not line up, instead of a model that loads clean and speaks gibberish.

9.3 Correct, without the download

The real checkpoint is gigabytes, and it is not needed to test the mapping. The mapping is pure name bookkeeping, and bookkeeping can be checked on a model small enough to build in a millisecond. Take a tiny forge model, write its weights out under HuggingFace’s names — the inverse of the table above — convert them back with forge’s real remap_from_hf, and load them into a fresh model through the very load_checkpoint guard from the last section. If the round trip is honest, every parameter returns exactly as it left and the guard passes in silence. If a single tensor lands under the wrong name or goes missing, that same guard raises.

Twenty-five parameters go out and twenty-five come back, every one bit-identical, and forge’s own guard signs off that no name went unaccounted for. That proves the translation table, the round trip, and the guard are internally consistent — every HuggingFace name maps to a real module and back, and the load refuses anything that does not. What it does not prove, and the chapter says so plainly, is the two things that need the real file: the gigabyte-scale read itself, and the conversion from the checkpoint’s bf16 storage to the working precision. Those ran once already, in Chapter 1, when the engine spoke. This check covers the logic that is cheap to break and cheap to test: the names.

The weights are in place. The engine now holds real knowledge, not random initialization, and can run on the trained model exactly as it did in Chapter 1. One black box is left between a string of text and the token ids the embedding expects: the piece that turns "The capital of France is" into the integers the model reads, and turns the integers it writes back into text. That is the next box: the tokenizer.