Final norm + logits
# grade this stage (from your course directory)
npx theinference test --stage final-norm-logits --watchThis is the step where your numbers meet OpenAI's. Apply the final LayerNorm, then project the last position against the tied token embedding to get next-token logits. Turn it green and your [50257] logits match the reference within 1e-3 — every score, for every token GPT-2 knows, produced entirely by code you wrote.
x_last[768] · wte[t][768] for t in 0..50256 → logits[t] (one dot product per vocab token)
= x_last @ wteᵀ where wteᵀ is [768,50257] → logits [50257]
What to build
Implement runLogits
Implement runLogits(inputIds): run the stack (reuse runHidden), apply ln_f (runLayerNorm with ln_f.{weight,bias}), then project the last position against the tied wte: logits[t] = x_last · wte[t]. Return the length-50257 vector.
export function runLogits(inputIds: number[]): number[] {
// -> logits for the final position, length 50257
}
Two mistakes get graded here: applying ln_f per block (it runs once, after block 11), and hunting for a separate LM-head matrix (there isn't one — dot x_last against the same wte rows you loaded back in Chapter 1). The quiz next step tests exactly these.
Your length-50257 logits are compared to the reference with allclose, atol 1e-3 — the loosest tolerance in the course. Float32-summation-order differences compound across 12 blocks of ~6 large matmuls each, so 1e-3 is the honest envelope of that drift, not slack.