Section 2 of 5 · LLM Foundations

Decoder Transformer Mechanics

The transformer should be a set of owned tensor operations rather than a diagram or framework class. This section connects shapes, causal information flow, modern block choices and parameter arithmetic.

7 units 24 topics 0/7

Units

7 in this section

2.1 Embeddings, positions and the vocabulary axis
explain
3 h
2.2 Scaled dot-product attention from scratch
derive · implement
5 h
2.3 Multi-head attention: split, merge and gradients
implement
4 h
2.4 MQA and GQA: sharing keys and values
derive · implement
4 h
2.5 RoPE: positions as rotations
derive · implement
4 h
2.6 The modern decoder block: residuals, RMSNorm and SwiGLU
explain · implement
6 h
2.7 Assembling and verifying a decoder
derive · implement
6 h

Before moving on

Checkpoint

Build a tiny decoder from blank tensor operations, derive every intermediate shape and parameter count, and pass causal intervention, mask, RoPE-offset, grouped-head, residual, normalization, tying and differential forward or backward tests.

Deliberately out of scope

Leave for later

MoE, MLA and broad modern-architecture catalogs
FlashAttention and custom GPU kernels
Long-context extension and production transformer libraries

Resources for this section

Shared by every unit above

Modern LLM Notebook - decoder construction path
Walking Labscourse

primary

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
Sebastian Raschkaimplementation

required

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
Stanford Universityassessment

required

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.
PyTorch scaled dot product attention reference
PyTorchreference

support

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
Deep-MLpractice

optional

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
Andrej Karpathyimplementation

support

Read model.py: MLP, Block, GPT.__init__'s special scaled initialization of c_proj.weight, and GPT._init_weights. Contrast the depth-dependent residual initialization with CS336's component initialization; do not adopt the training or generation code.

Apply the section

Mini-project · Tiny decoder with an independent oracle (10 to 14 h)

Build a configurable decoder-only model from tensor operations in the same repository, and prove it correct without trusting a training curve: public snapshot tests for each component, a slow per-head reference for attention, and intervention tests for causality.

Steps

Implement linear, embedding, RMSNorm, SwiGLU, RoPE, softmax and scaled dot-product attention from tensor operations, and pass the matching CS336 tests in tests/test_model.py and tests/test_nn_utils.py.
Add grouped KV heads to your attention with an explicit map, query head \(h\) reads KV head \(\lfloor h/(H_q/H_{\mathrm{kv}}) \rfloor\), so that MHA and MQA are the two ends of one code path.
Write a slow reference that loops over heads with sliced weights, and compare outputs and gradients in float64 with your batched path for MHA, GQA and MQA.
Assemble the decoder from one config object with optional embedding tying, and pass test_transformer_block, test_transformer_lm and test_transformer_lm_truncated_input.
Derive the parameter count of your config by hand and assert that it equals the count of unique registered parameters, tied and untied.
Write property tests showing RoPE preserves norms with scores depending only on relative position including offset equivalence, a future-token change leaves earlier outputs unchanged, and batch rows stay independent.

Deliverables

The model package: components, config, grouped attention and the tied or untied head.
Your test suite: the head-loop oracle, float64 gradient comparisons, suffix-intervention tests and the parameter-count assertion.
A one-page shape and parameter sheet for your config, derived by hand.

Checks

external-tests · CS336 `tests/test_model.py` and the softmax test
rubric · Self-grade with the rubric below

Rubric

Criterion	Pass	Strong
Component correctness	All CS336 model tests and the softmax test pass.	Your own tests add fully masked rows, padding plus causal masks and non-contiguous inputs.
Independent oracle	The batched attention matches the head loop on outputs.	Outputs and all projection gradients match in `float64` for MHA, GQA and MQA within a stated tolerance.
Causality	Changing a suffix leaves prefix logits unchanged.	Your tests also show that the change does reach the logits at and after the edited position, and they cover batch independence.
Grouped heads and RoPE	Your batched attention matches the head-loop reference for MHA, GQA and MQA, and the RoPE, causality and batch-independence tests pass.	You explain why an interleaved KV expansion differs from the PyTorch grouping convention, and your weights load correctly under yours.
Accounting	The hand-derived parameter count equals the registered count.	The sheet also gives the dominant forward FLOP terms and predicts how each scales when the sequence length doubles.

Feeds into: Capstone artifact 2, the tiny decoder stack. Section 3 trains this model, and section 5 adds a cache to it.