LLM Developer Docs

This is a small language model built from scratch. This documentation shows the mathematics behind it, the code that implements the mathematics, and how the pieces work together.

Start learning Explore the code

The model at a glance

Read from config.py at the pinned commit.

Setting	Name in `config.py`	Value
Model
Layers	num_layers	`2`
Heads	num_heads	`2`
Width	emb_dim	`32`
Head width No module reads this constant.	head_dim	`16`
FFN width	ff_dim	`128`
Window length	seq_len	`16`
Context length	context_len	`32`
Tokenizer
BPE merges	num_merges	`500`
Optimizer
Largest learning rate	lr_max	`1e-3`
Smallest learning rate	lr_min	`1e-4`
Warmup steps	warmup_steps	`50`
Beta 1	beta1	`0.9`
Beta 2	beta2	`0.999`
Adam epsilon	adam_eps	`1e-8`
Weight decay	weight_decay	`0.01`
Gradient clip	grad_clip	`1.0`
Sampling defaults
Temperature	temperature	`0.8`
Top-k 0 turns it off.	top_k	`0`
Top-p	top_p	`0.9`
New tokens in talk.py	max_new_tokens	`50`

The pipeline

Every step the code runs, from text to generated tokens. Shapes are for one sequence: the model has no batch dimension.

Text to training windows

Embeddings

5. Token plus position embeddings(T, d)(16, 32)transformer.py:158

Blocks, repeated num_layers = 2 times transformer.py:162

6. LayerNorm 1(T, d)(16, 32)transformer.py:133
7. Attention(T, d)(16, 32)transformer.py:134
8. Residual add(T, d)(16, 32)transformer.py:135
9. LayerNorm 2(T, d)(16, 32)transformer.py:136
10. Feed-forward(T, d)(16, 32)transformer.py:137
11. Residual add(T, d)(16, 32)transformer.py:139

Output

12. Final LayerNorm(T, d)(16, 32)transformer.py:165
13. Linear head(T, V)(16, V)transformer.py:166
14. Logits(T, V)(16, V)transformer.py:168

Loss and update

Generation

19. Generation(V,)Samples from the last row of logits.generate.py:33

Plain names have no page yet. The link under each name opens the real lines at the pinned commit.

Labels used on this site

Illustrative Example: Hand-picked numbers. These are not learned values.
Actual Model Output: Numbers produced by running the repository, with their source stated.
Derived from implementation: Mathematics the code implements but the reference PDF does not state.

The model at a glance

The pipeline

Questions this site answers

Labels used on this site