Build your own GPT-2
Load the GPT-2 weights·8 min

Parse the header

0 / 1
# grade this stage (from your course directory)
npx theinference test --stage parse-the-header --watch

You have the map; now read the file and make it talk back. Build a loader that parses the header and pulls a couple of tensors out by name — then print a small summary so you (and the grader) can see the real numbers fall out of the bytes.

What to build

Your task

Print the weight summary

When run as ./your_program.sh load-weights <path-to-model.safetensors>, print a single JSON object to stdout matching the shape below.

  • numTensors — how many entries are in the JSON header (ignore any __metadata__ key).
  • wteShape / wpeShape — the shapes of wte.weight (token embeddings) and wpe.weight (position embeddings).
  • wteSample — the first five F32 values of the wte.weight tensor.
json
{
  "numTensors": 160,
  "wteShape": [50257, 768],
  "wpeShape": [1024, 768],
  "wteSample": [<wte[0][0]>, <wte[0][1]>, <wte[0][2]>, <wte[0][3]>, <wte[0][4]>]
}
Exactly what the reference prints
{
  "numTensors": 160,
  "wteShape": [50257, 768],
  "wpeShape": [1024, 768],
  "wteSample": [-0.1101, -0.0393, 0.0331, 0.1338, -0.0485]
}

Sanity-check against this before you trust anything downstream: if your numTensors isn't 160, or your wteSample doesn't start with -0.1101, your offset or endianness math is off. Fix it here — nothing in the rest of the course can work if the very first floats are wrong.

Read it with a DataView

Two reads do the whole job:

  • Header length: view.getBigUint64(0, true) — the true is little-endian. It returns a BigInt, so Number(n) it before using it as an offset.
  • Values: new Float32Array(buf, 8 + N + begin, count) views the tensor's bytes directly as floats (or view.getFloat32(off, true) one at a time).

That true flag is the #1 silent bug in this step: leave it off and the reads default to big-endian, so every number comes back as scrambled nonsense with no error thrown.

160, not 148 — and don't "fix" it

The header has 160 tensors. But 12 of them are vestigial h.{i}.attn.bias causal-mask buffers of shape [1,1,1024,1024] — leftover precomputed masks GPT-2 doesn't need at inference (you'll build your own mask in Chapter 2). That leaves 160 − 12 = 148 "real" weights.

The summary still wants 160: count every header entry, skipping only the __metadata__ key. Don't try to filter the mask buffers out to "clean up" the count — that's a wrong answer, not a fix.

Read the path from argv

Your program is run with the path to the weights; the JSON it prints is compared to the reference within atol 1e-4. The path is provided by the grader — read it from argv, don't hardcode it.

The safetensors formatsafetensors vs pickle