Week 8. Assembling the full transformer

Phase 2. The modern transformer · week 8 of 24

Learn in the app: tutor, coding problems →

Core: the full pass and the loss, initialization, matching logits against a real model, a transformer from an empty file · Depth: track D and the second half of track C · ≈ 11 h core / 19 h total

The pieces from weeks 6–7 become a model, and the only real test appears: matching the logits (the model's raw outputs before softmax) against someone else's reference implementation. Everything that follows (training, inference, sampling) silently assumes the model is correct. A bug here does not crash; it quietly degrades everything built on top.

Step 1. The full pass and the loss

The full forward: embed → L×(RMSNorm → attn → residual → RMSNorm → FFN → residual) → norm → head. Shapes: ids (B, T) → embeddings (B, T, D) → the blocks do not change the shape → final norm → logits (B, T, V)
The loss with a shift: position t predicts token t+1, i.e. logits[:, :-1] against targets[:, 1:]. A sanity check at the start: an untrained model gives a nearly uniform distribution, so the loss ≈ ln V (with V = 32 000 that is 10.37). If it is much higher, the initialization scale is broken

Step 2. Initialization

In plain terms. In the residual stream, each of the 2L sublayers adds its own contribution. If the contributions are independent and each has variance 1, the variances add up. With L = 8 that is 16 additions: the variance of the stream is 1 + 16 = 17, and the spread has grown by a factor of √17 ≈ 4. The deeper the model, the stronger the blow-up. If each contribution is scaled down by √(2L), the sum of all additions has variance 1 at any depth.

Initialization, scaled init for residual branches: all matrices are N(0, 0.02²), while the output projections of the branches (w_o, w_down) get std = 0.02/√(2L). Derivation (D7): Var(x_L) = Var(x_0) + 2L·σ², and σ² ∝ 1/(2L) makes the sum independent of depth
Dropout: why modern LLMs barely use it. Dropout fights memorization of repeated examples, but in pretraining the model sees almost every token only once. There is nothing to memorize, and the noise slows convergence. It is sometimes brought back for fine-tuning on small datasets

Step 3. Matching logits against a real model

In plain terms. When two ledgers show different totals, you do not recount them from scratch. You go page by page looking for the first one where the sums diverge: everything before it is correct, and the error is on it. Matching a model works the same way: compare the activations after the embeddings, after each block, inside the block, and look for the first place where the discrepancy exceeds rounding noise.

The procedure: fp32, identical input ids, forward hooks on every block, tolerance ~1e-4
Where it breaks most often. The RoPE in HF implementations of Llama/Qwen rotates pairs (i, i + H/2), while nanolm rotates adjacent (2i, 2i+1): the rows of W_q and W_k within each head must be permuted. In nanolm Swish sits on w_up, while in HF it is on gate_proj: so gate_proj → w_up, up_proj → w_gate. Then check the norm ε, rope_theta, whether the embeddings are tied, and an explicit F instead of 8D/3

Common mistakes

Losing the one-token shift: the model learns to copy its input, and the loss drops to zero suspiciously fast. There is no dedicated test. Leaking the future through the mask is caught by test_model_is_causal
Copying the matrix instead of tying it: two tensors instead of one, caught by test_weight_tying_shares_one_tensor
A layer that is created but not connected to the graph: test_gradients_reach_every_parameter. The general health check is test_model_can_overfit_one_batch: if the model cannot memorize a single batch, do not go any further

Code → nanolm/model.py: the full model: Block, NanoLM (_init_weights, scaled init in __init__, forward(input_ids, targets, cache) with the shift, num_parameters). Load the weights of a real small model (for example, Qwen3-0.6B) into your implementation and match the logits. This is the only real test. Qwen3 has QK-norm (Attention supports it, but Block does not enable it), and its head dimension is set separately: N·H ≠ D. It is easier to start with a Llama-architecture model, where H = D/N and there is no QK-norm. The tests in tests/test_model.py import the reference directly: for your own model, copy them and change the import.

Track E starts here and runs to the end of the program: every Friday you write a transformer from an empty file against the clock. The goal: get from ~90 minutes down to ~25. The analysis of 57 interviews that the program was checked against says it plainly: implementing and debugging a transformer comes up in interviews so often that you simply cannot afford to lose points on it.

Math (track D): D7: variance growth in the residual stream and choosing the initialization scale; D8: the rank of the attention matrix, and why softmax raises the rank of QKᵀ while linear attention does not.

Interview question of the week: "Your logits differ from the reference by 1e-2. How will you track it down?" A 3-minute structure: the same ids and tokenizer → fp32 → hooks, the first diverging layer → inside the layer: the norm (ε, fp32), q and k after RoPE (the pairing convention), the mask, the repeat_kv order, the SwiGLU gate → restore the original precision and state the expected tolerance.

Sources: Radford et al., GPT-2 (2019), scaled initialization; Touvron et al., LLaMA (2023).

Week outcomes

I can write the full forward pass of a transformer and load a real model's weights into it with matching logits.
I can find the cause of a logit mismatch by comparing activations layer by layer.
I can explain why the initialization of the residual branches' output projections is scaled down as depth grows.
I can write a transformer from an empty file and record my starting time for track E.

Self-check

List, in order, every operation of the forward pass with the tensor shape after each one.
What happens to the variance of the residual stream after L layers without scaled init?
Why do modern LLMs barely use dropout?