Unit 2.7 · about 6 h
303 / 324Assembling and verifying a decoder
Correct components do not automatically make a correct model. Assembly checks catch final-norm, configuration, weight-sharing and information-flow defects before training hides them.
- derive
- implement
You can now
- Assemble a configurable decoder from owned tensor components, ending in raw \([B, T, V]\) logits.
- Derive component parameter counts and dominant forward matrix-multiplication terms, distinguishing tied and untied heads.
- Verify causal prefix equivalence and match an independent reference under transferred weights.
- Contrast the decoder's causal next-token use with BERT's bidirectional masked-token use.
Covers
- Embedding-output weight tying and final logits
- Decoder composition and configuration ownership
- Parameter counts and dominant operation terms
- Slow attention oracle and causal intervention tests
- Encoder-only BERT contrast
Study, in this order
4 steps
- Decoder assembly, configuration and tyingModern LLM Notebook - decoder construction path
Implementing
MiniGPTfrom scratch: Implementation ofMiniGPTand Logits and predictions; Reading the nanoGPT source code: The main thread ofGPT.forwardand Weight Tying: sharing weights between input and outputDepth · Read architecture and shape flow only. Skip training, generation demos and the full GPT-2 parameter survey; substitute your modern blocks for the teaching version.
- Full tensor model and tied parameter countLLMs from Scratch companion - attention and GPT construction
Section 4.6, Coding the GPT model, through the weight-tying and memory-count discussion
Depth · Skip the larger GPT-2 configuration exercise and section 4.7 generation. Keep the final normalization and raw-logit contract.
- Full model and resource accountingStanford CS336 Assignment 1 - model contract
Section 3.1, Transformer LM, printed pp. 13 to 14; section 3.5, The Full Transformer LM, including (
transformer_lm), Resource accounting and (transformer_accounting), printed pp. 26 to 28Depth · Do
transformer_accountingparts (a) to (c). Leave the model-size sweep and long-context sweep in parts (d) and (e) for the model-arithmetic section. - Optional encoder-only contrastModern LLM Notebook - decoder construction path
Decoder-Only Architecture: the Encoder-Only, Encoder-Decoder and next-token-training comparison
Depth · Use only for the optional BERT contrast. Do not start an encoder implementation or model survey.
Focus on
- One configuration owns \(V\), \(D\), layer count \(L\), query and KV head counts, FFN width \(F\), RoPE \(\theta\), context policy and tying. Validate divisibility before creating weights; different layers share a design, not all their parameters.
- A tied head uses the same \([V, D]\)
Parameteras input lookup and computes \(h @ E^\top\). Equal initial values in separate Parameters are not tying; gradients from lookup and classification must accumulate into one weight. - For bias-free MHA with two RMSNorms and SwiGLU per block, untied parameters are \(2VD + L(4D^2 + 3DF + 2D) + D\). Tying subtracts \(VD\); standard RoPE adds no learned parameters.
- List dense forward terms: per layer \(8BTD^2\) for MHA projections, \(4BT^2D\) for \(QK\) and \(AV\), and \(6BTDF\) for SwiGLU matrices; the vocabulary head adds \(2BTDV\). Use the slow per-head oracle and suffix interventions to check causality, not a loss curve.
Reproduce closed-book
Traps
- Adding softmax inside the model's training forward path. Return raw logits; the loss owns normalization.
- Matching a CS336 snapshot with a tied output head. Its supplied embedding and
lm_headweights are independent; test untied mode there and tying separately. - Checking causality only in the attention matrix. Cross-token normalization or a composition error can still alter prefix logits.
Practice
CS336 Assignment 1: `transformer_lm`
Problem (transformer_lm), section 3.5; implement run_transformer_lm and pass test_transformer_lm and test_transformer_lm_truncated_input
Use the untied MHA contract for snapshots. Add prefix-alone versus full-sequence comparisons, adversarial suffix changes, tied-weight identity and gradient accumulation checks, and a full-model comparison using the slow attention oracle. Keep dropout off.
CS336 Assignment 1: `transformer_accounting`, parts (a) to (c)
Problem (transformer_accounting), printed pp. 27 to 28: parameter and FP32 weight memory count, forward matrix multiplies and dominant FLOPs
This written problem has no public accounting test. Check its algebra against your tiny model's unique parameter count, using tests/test_model.py's model tests for implementation correctness; do not allocate the XL model.
Self-test
-
Does tying the head remove the vocabulary projection's compute?
No. It removes a separate \(VD\) weight matrix, but \(h @ E^\top\) still costs about \(2BTDV\) dense FLOPs. -
What should changing tokens after zero-based position \(i\) do to logits through \(i\)?
Nothing within numerical tolerance, with identical prefix positions and deterministic evaluation. Those logits may depend on tokens at positions at most \(i\), not the suffix. -
If sequence length doubles at fixed batch and width, which dominant terms quadruple?
Dense \(QK\) and attention-times-\(V\) work, and the stored score/probability matrices, scale quadratically. Projection and FFN terms scale linearly with length. -
Why cannot an ordinary BERT forward pass serve as this causal decoder oracle?
BERT uses bidirectional attention for masked-token prediction, so its representations can use right context. The decoder must preserve the causal prefix boundary.
Completion criterion
Stop when
Both full-model public tests pass, the tiny parameter count matches the derivation, and your model matches the slow reference and preserves prefix logits under suffix replacement in tied and untied modes.
Mark passed, then open 3.1 Only mark this when you can do it without notes.Leads to
- 3.1 The next-token objective: shifts, masks and perplexity Turn raw vocabulary logits into the shifted next-token loss.
- 4.6 Counting parameters and FLOPs Extend the component count into a full model arithmetic report.
- 5.3 Proving the cache correct Reuse prefix equivalence when validating incremental execution.
Comes after
- 2.4 MQA and GQA: sharing keys and valuesMake the head structure an explicit model configuration choice.
- 2.6 The modern decoder block: residuals, RMSNorm and SwiGLUStack these blocks and add the final norm and vocabulary head.
Resources for this unit
- Modern LLM Notebook - decoder construction path Work through English notebooks 03-embedding, 04-position-encoding, 05-transformer-block, 06-mini-gpt and the selected normalization, RoPE, GQA and SwiGLU comparisons in 08-gpt2-to-modern-models. Re-derive shapes and rerun with changed dimensions.
- LLMs from Scratch companion - attention and GPT construction Use the ch03 attention notebooks and ch04 model code. Rebuild the model without a high-level Transformer module and compare one forward and backward path with an independent small reference.
- Stanford CS336 Assignment 1 - model contract Use the named linear, embedding, rmsnorm, positionwise_feedforward, rope, softmax, scaled_dot_product_attention, multihead_self_attention, transformer_block, transformer_lm and transformer_accounting specifications and public tests. Keep implementation independent and respect the course AI and honor-code boundary.
Also in this section
- PyTorch scaled dot product attention reference Use the functional contract only after the explicit attention oracle passes. Transfer identical weights and masks, compare outputs and gradients, and explain dtype, causal and dropout behavior from the documented semantics.
- Deep-ML candidate - three transformer closing drills After the platform quality sample passes, attempt at most problems 490, 53 and 904 as changed-prompt drills for attention shapes and multi-head composition. Preserve independent tests; a judge result does not replace the section checkpoint.
- nanoGPT - residual projection initialization Read
model.py:MLP,Block,GPT.__init__'s special scaled initialization ofc_proj.weight, andGPT._init_weights. Contrast the depth-dependent residual initialization with CS336's component initialization; do not adopt the training or generation code.