9. Real weights: safetensors and GPT-2

Why load, and what you’re loading

Chapter 7’s estimate put Qwen3-0.6B’s pretraining at about $10^{23}$ FLOPs. Nobody reproduces that to run a model. Open-weight models publish the result of training as files, and an inference engine’s first job is to read those files into exactly the computation that produced them.

A Hugging Face model snapshot is a folder:

models/gpt2/
├── config.json              architecture hyperparameters: layers, widths, vocab size, epsilon, ...
├── tokenizer.json           the tokenizer: vocabulary, merges, special tokens, pre-tokenizer
├── tokenizer_config.json    chat template, special-token roles
├── generation_config.json   default sampling settings and stop tokens
└── model.safetensors        the weights (large models: model-00001-of-00004.safetensors ... + an index)

Download GPT-2 small (548 MB) to follow along:

uv pip install -r optional-requirements.txt
hf download openai-community/gpt2 --local-dir models/gpt2

The safetensors format

Older checkpoints were Python pickles (pytorch_model.bin). Unpickling can execute arbitrary code, so loading an untrusted model was a security hole. They also had to be fully deserialized before you could look at them. Safetensors fixes both with a deliberately simple layout:

┌───────────────┬──────────────────────────────┬────────────────────────────────────────┐
│ 8 bytes: N    │ N bytes: JSON header (UTF-8) │ tensor data: raw little-endian bytes   │
│ (u64, LE)     │                              │ at the offsets the header lists        │
└───────────────┴──────────────────────────────┴────────────────────────────────────────┘

The header maps each tensor name to its dtype, shape and byte range within the data section:

{
  "__metadata__": {"format": "pt"},
  "wte.weight": {"dtype": "F32", "shape": [50257, 768], "data_offsets": [0, 154389504]},
  "h.0.attn.c_attn.weight": {"dtype": "F32", "shape": [768, 2304], "data_offsets": [...]}
}

That’s the entire format. Because data is raw bytes at known offsets, a reader can memory-map the file and read any tensor without touching the others. A loader can inspect a 300 GB checkpoint’s header in milliseconds, which you’ll do in Chapter 30. You can write a valid file by hand:

import json, struct
header = json.dumps({"w": {"dtype": "F32", "shape": [2], "data_offsets": [0, 8]}}).encode()
blob = struct.pack("<Q", len(header)) + header + struct.pack("<2f", 1.5, -2.0)
open("hand.safetensors", "wb").write(blob)        # a real, loadable checkpoint with one tensor

The reader is just as short:

def read_header(path):
    """Return (header dict without __metadata__, metadata dict, byte offset of the data section).  (Your engine: Chapter 9)"""
    with open(path, "rb") as file:
        (length,) = struct.unpack("<Q", file.read(8))
        if length > 100_000_000:
            raise ValueError("Implausibly large safetensors header")
        header = json.loads(file.read(length))
    metadata = header.pop("__metadata__", {}) or {}
    return header, metadata, 8 + length

def load_file(path, device="cpu", names=None):
    """Read tensors into a dict. Memory-maps the file and copies one tensor at a time.  (Your engine: Chapter 9)"""
    header, _, data_start = read_header(path)
    selected = header if names is None else {n: header[n] for n in names}
    tensors = {}
    with open(path, "rb") as file, mmap.mmap(file.fileno(), 0, access=mmap.ACCESS_READ) as view:
        for name, info in selected.items():
            dtype = DTYPES[info["dtype"]]
            begin, end = info["data_offsets"]
            count = (end - begin) // torch.empty((), dtype=dtype).element_size()
            if count != _numel(info["shape"]):
                raise ValueError(f"{name}: byte range does not match its shape")
            if count == 0:                                   # frombuffer rejects empty reads
                tensors[name] = torch.empty(info["shape"], dtype=dtype, device=device)
                continue
            with warnings.catch_warnings():
                warnings.simplefilter("ignore")          # frombuffer warns that mmap is read-only
                flat = torch.frombuffer(view, dtype=dtype, count=count, offset=data_start + begin)
            tensors[name] = flat.reshape(info["shape"]).to(device, copy=True)
    return tensors

torch.frombuffer creates a tensor that views the mapped bytes with the right dtype, and .to(device, copy=True) makes an owned copy, on the GPU if requested. One tensor at a time keeps peak host memory low. Large models split their weights across shards plus a model.safetensors.index.json that maps each tensor to its file; snapshot_tensors walks them and checks that the index and files agree.

Dtypes, and BF16 for free

Header dtypes are short codes: F32, F16, BF16, I64, U8, F8_E4M3 and so on. BF16 is the most common for modern checkpoints, and it’s worth knowing why it’s so convenient. A BF16 number is exactly the top 16 bits of an FP32: same sign, same 8-bit exponent, mantissa cut from 23 bits to 7. Widening BF16 to FP32 is a 16-bit shift. FP16 has a different exponent width and needs real conversion logic:

# PyTorch does this for you: torch.frombuffer(raw, dtype=torch.bfloat16).float()
# By hand, on the raw 16-bit integers:
def bf16_to_f32(bits16):            # int16 tensor of raw BF16 bit patterns
    return (bits16.to(torch.int32) << 16).view(torch.float32)

// Read just the header: 8-byte little-endian length, then that many bytes of JSON.
inline std::string safetensors_header(const std::string& path) {
    std::ifstream f(path, std::ios::binary);
    if (!f) throw std::runtime_error("cannot open " + path);
    uint64_t n = 0;
    f.read(reinterpret_cast<char*>(&n), 8);           // assumes a little-endian host, like x86 and ARM
    std::string json(n, '\0');
    f.read(json.data(), std::streamsize(n));
    return json;
}
inline float bf16_to_float(uint16_t bits) {
    uint32_t widened = uint32_t(bits) << 16;           // BF16 is the top half of an FP32
    float out;
    std::memcpy(&out, &widened, 4);
    return out;
}

#![allow(unused)]
fn main() {
/// BF16 is the top 16 bits of an f32: widening is a shift.
#[inline]
pub fn bf16_to_f32(bits: u16) -> f32 {
    f32::from_bits((bits as u32) << 16)
}

/// IEEE half precision: 1 sign, 5 exponent, 10 mantissa bits.
pub fn f16_to_f32(bits: u16) -> f32 {
    let sign = if bits >> 15 == 1 { -1.0 } else { 1.0 };
    let exponent = ((bits >> 10) & 0x1f) as i32;
    let mantissa = (bits & 0x3ff) as f32;
    match exponent {
        0 => sign * mantissa * 2f32.powi(-24),                       // subnormal
        31 => if mantissa == 0.0 { sign * f32::INFINITY } else { f32::NAN },
        e => sign * (1.0 + mantissa / 1024.0) * 2f32.powi(e - 15),
    }
}
}

The Rust engine keeps weights in BF16 in memory and widens each value inside the matrix-vector product. That’s half the bytes of FP32, so decode runs about twice as fast, for exactly the reason Chapter 1 gave.

Mapping a checkpoint onto your model

Loading means: for every tensor in the file, find the parameter in your model it belongs to, check its shape, and copy it. Every tensor must be used, and every parameter must be filled. Reject the load if either set has leftovers, because a silently uninitialized parameter produces plausible-looking garbage.

GPT-2’s names map onto your Chapter 6 attribute names like this:

The transpose trap

GPT-2 was written with a Conv1D layer that stores weights as [in, out], while nn.Linear stores [out, in] (Chapter 2). Three of the four matrices change shape when transposed, so forgetting the transpose fails loudly. But c_proj in attention is [768, 768]: square. Copying it untransposed passes the shape check and computes $x W$ instead of $x W^\top$. The model then produces fluent-looking nonsense. This is the canonical example of why shape checks are necessary but not sufficient.

@torch.no_grad()
def load_gpt2(directory, device="cpu", dtype=torch.float32):
    """GPT-2 from Hugging Face's safetensors layout.  (Your engine: Chapter 9)

    GPT-2 stores its projections in a Conv1D layout [in, out]; nn.Linear wants [out, in],
    so those four weights are transposed. The head is tied to the token embedding.
    """
    raw = read_config(directory)
    if raw.get("model_type") != "gpt2" or raw.get("activation_function", "gelu_new") != "gelu_new":
        raise ValueError("Expected a standard GPT-2 checkpoint")
    model = GPT(GPTConfig(vocab=raw["vocab_size"], context=raw["n_positions"], width=raw["n_embd"],
                          layers=raw["n_layer"], heads=raw["n_head"],
                          eps=raw.get("layer_norm_epsilon", 1e-5))).to(device=device, dtype=dtype)
    targets = {"wte.weight": (model.token.weight, False), "wpe.weight": (model.position.weight, False),
               "ln_f.weight": (model.norm.weight, False), "ln_f.bias": (model.norm.bias, False)}
    for i, block in enumerate(model.blocks):
        for source, module, transpose in [("ln_1", block.ln1, False), ("ln_2", block.ln2, False),
                                          ("attn.c_attn", block.attn.qkv, True),
                                          ("attn.c_proj", block.attn.proj, True),
                                          ("mlp.c_fc", block.up, True), ("mlp.c_proj", block.down, True)]:
            targets[f"h.{i}.{source}.weight"] = (module.weight, transpose)
            targets[f"h.{i}.{source}.bias"] = (module.bias, False)
    loaded, head_alias = set(), None
    for name, value in snapshot_tensors(directory):
        name = name.removeprefix("transformer.")
        if name.endswith((".attn.bias", ".attn.masked_bias")):   # old causal-mask buffers, not weights
            continue
        if name == "lm_head.weight":                               # an alias of wte when present
            head_alias = value
            continue
        if name not in targets:
            raise ValueError(f"Unexpected tensor {name}")
        parameter, transpose = targets[name]
        assign(parameter, value.T if transpose else value, name)
        loaded.add(name)
    if set(targets) - loaded:
        raise ValueError(f"Missing tensors: {sorted(set(targets) - loaded)[:5]}")
    if head_alias is not None and not torch.equal(head_alias.to(model.token.weight), model.token.weight):
        raise ValueError("lm_head.weight disagrees with the tied token embedding")
    return model.eval()

The packed c_attn also fixes the order of Q, K and V within its 2,304 outputs (Q first, then K, then V), which matches the chunk(3) in your attention. A checkpoint that packed them as K, Q, V would need reordering, and nothing in the shapes would tell you.

Proving the load is correct

“It loaded without errors” proves little. Correctness has levels, and you should climb them in order:

1. Shapes and names: every tensor used, every parameter filled. The loader enforces this.
2. Logits match an independent implementation on the same input IDs. Hugging Face Transformers is the usual reference. With FP32 on both sides, expect agreement to about 1e-4 or better. Your milestone test builds a small random GPT-2 with Transformers, saves it, loads it with your code, and compares logits to 1e-4.
3. Greedy generations match for several prompts.

When logits disagree, don’t stare at the final output. Find the first intermediate that differs. Run both models with hooks that record each stage, and compare in order: token embeddings, position embeddings, first LayerNorm, Q/K/V of block 0, attention output, MLP output, the stream after block 0, and so on. The first mismatch localizes the bug to one operation, almost always an orientation, an epsilon, an activation variant, or an axis.

Tip

Compare fixed inputs, not generated text. Once two models choose different tokens, everything after is different input, and the comparison is meaningless.

Run GPT-2 in your engine

python run.py gpt2 --model-dir models/gpt2 --prompt "The GPU is a" --new-tokens 30

The tokenizer (from transformers) handles text at the boundary; everything else runs in your code: your safetensors reader, your loader, your GPT, your sampler. GPT-2 small is a 2019 model, so expect grammatical but meandering continuations, which is a big step up from Chapter 7’s model.

Build it

pytest tests/test_ch09_weights.py
python run.py gpt2 --impl engine --model-dir models/gpt2

Tip

torch.frombuffer(view, dtype=..., count=..., offset=...) reads from a memory map without copying. It rejects zero-length reads, so handle empty tensors separately. Silence its “non-writable buffer” warning, then copy (.to(device, copy=True)) so the tensor outlives the mapping.

Stretch exercises

1. ★ Write inventory(dir) output for GPT-2: tensor count and total bytes. Does the total match the parameter formula from Chapter 6 times 4 bytes? Why are there extra tensors? Where: experiments/ch09.py (create it), calling the provided engine.safetensors_io.inventory.
2. ★★ Load GPT-2 in BF16 instead of FP32 and measure the maximum logit difference against the FP32 load. Is greedy text the same for a 50-token generation? Where: experiments/ch09.py (create it), calling engine.loaders.load_gpt2 with each dtype.
3. ★★ Deliberately skip the transpose for c_proj only. Which tests or checks catch it? What does the generated text look like? Where: temporarily change the c_proj assignments in load_gpt2 in engine/loaders.py, then restore them.
4. ★★★ Write a converter that saves your Chapter 7 model as a Hugging Face–compatible GPT-2 snapshot (config, transposed weights, names), then load it with transformers and compare logits. Where: new experiments/export_gpt2.py, using engine.train.load_checkpoint and engine.safetensors_io.save_file.

Check your understanding

1. Why can a safetensors header be read without loading any tensor data?
2. Why is loading a pickle-based checkpoint from an untrusted source dangerous?
3. Why can a missing transpose on a square matrix pass a shape check?
4. A tied checkpoint omits lm_head.weight. Why is that valid, and what must the loader do?
5. Your logits differ from the reference. What’s the first thing to compare?

Going deeper

• BALLM §§5.4-5.5 (pp. 159-168): saving and loading weights, and loading OpenAI’s GPT-2 into a from-scratch model (from TensorFlow checkpoints, with the same transposes).
• The safetensors specification and its security audit.
• Hugging Face’s modeling_gpt2.py, especially Conv1D, the origin of the [in, out] layout.