Module 3 · Complete reference implementation

Build the TransformerModel, piece by piece.

This runnable PyTorch model joins the code checkpoints from Lessons 15–23 into one decoder-only Transformer that maps token IDs to next-token logits.

About 180 linesPyTorch onlyDesigned for tracing

Implementation map

Read down the table once, then follow the source in the same order. Every row corresponds to one conceptual checkpoint in the course.

Lesson	Code element	Shape contract
15 · Hidden states	`TransformerConfig`, `token_embedding`	\([B,T]\rightarrow[B,T,D]\)
16 · Normalization	`RMSNorm`	\([B,T,D]\rightarrow[B,T,D]\)
17 · Q/K/V and position	`split_heads`, `precompute_rope`, `apply_rope`	\([B,T,D]\rightarrow[B,H,T,D_h]\)
18 · Causal mask	scaled scores and `masked_fill`	\([B,H,T,D_h]\rightarrow[B,H,T,T]\)
19 · Weighted values	`softmax(scores) @ v`	\([B,H,T,T]\rightarrow[B,H,T,D_h]\)
20 · Join heads	`transpose`, `view`, `out_proj`	\([B,H,T,D_h]\rightarrow[B,T,D]\)
21 · Residual block	`TransformerBlock`	\([B,T,D]\rightarrow[B,T,D]\)
22 · SwiGLU	`SwiGLU`	\(D\rightarrow D_{\text{ff}}\rightarrow D\)
23 · Full model	`TransformerModel`, `lm_head`	\([B,T,D]\rightarrow[B,T,V]\)

Retained from MiniMind

Decoder-only causal attention, RMSNorm, rotary position encoding on Q/K, pre-norm residual paths, SwiGLU, stacked blocks, final normalization, and tied token/output embeddings.

Deliberately omitted

Grouped-query attention, KV caching, Flash Attention, YaRN scaling, dropout, mixture-of-experts routing, Hugging Face wrappers, generation sampling, and training loss.

Reference lineage: teaching adaptation of MiniMind model_minimind.py at commit d65ef2c. Names and tensor layout were simplified to match this course's notation.

Run the model

Download transformer_model.py, install PyTorch, then run:

python transformer_model.py

# Expected shape lines:
# input: (1, 4)
# logits: (1, 4, 32)

Read the output: there is one vector of 32 vocabulary logits for each of the four input positions. A tokenizer and sampling loop would turn those logits into generated text; they are outside this core architecture lesson.

Complete source

First locate each class from the implementation map. Then trace one tensor through TransformerModel.forward, one TransformerBlock, and CausalSelfAttention.forward.

Loading transformer_model.py…

Code-reading exercises

Change hidden_size to 18 while keeping 4 heads. Which validation fails, and why?
For the example configuration, calculate head_dim before running the file.
Identify the one line that prevents future information from changing an earlier token representation.
Temporarily remove out_proj. Which conceptual role is lost even if the outer shape still matches?
Explain why lm_head changes the final axis from \(D\) to vocabulary size \(V\).
Add a forward hook or print statement that records the shape after every block.

Check guidance

hidden_size % num_heads fails because heads cannot receive equal feature widths.
\(16/4=4\).
The upper-triangular mask followed by masked_fill(..., -inf).
The learned mixing of concatenated head features back into model coordinates.
Each model feature vector must become one score for every possible next token.
Print x.shape inside the layer loop or register a hook on each TransformerBlock.