This project was made to demonstrate that I understand what's inside an LLM, not that I can call one. Every core piece is written by hand rather than imported.
A decoder-only transformer built and trained end-to-end, with every core piece written by hand instead of imported. The attention math, rotary positional embeddings, layer norm, weight init, training loop, and even the Wikipedia-markup stripper are hand-rolled.
Model
- 44.6M params (18.9M non-embedding), 6 layers, 8 heads, d=512
- 512-token context length
- GPT-2 BPE tokenizer via
tiktoken(the one external component, see trade-offs)
Training
- ~0.92B tokens, streamed directly from official Wikipedia dumps
- Music-related articles upweighted 3× (Music and instruments are a major hobby of mine, so I wanted it this model to be good at music)
- One consumer GPU target; this run used Colab Pro (a single L4)
Dependencies
torch,numpy,tiktoken,requests.
- Results
- What broke
- Quickstart
- Repo layout
- Architecture, and why
- The data: Wikipedia, with a music bias
- What I'd do differently at scale
- References
model.py the transformer: LayerNorm, RoPE, causal attention, blocks, generation
config.py model presets (debug/small/base) + training hyperparameters
data.py Wikipedia dump -> plain text -> GPT-2 BPE -> train.bin/val.bin
train.py training loop: AMP, cosine LR, grad accumulation, eval, checkpoints, samples
sample.py generate from a checkpoint (temperature / top-k / top-p)
plot_loss.py metrics.csv -> loss_curve.png
tests/ correctness tests for the things that break silently
out/ plotting & training output
Final run: 13,999 steps x 65,536 tokens/step ~= 0.92B tokens (the ~20
tokens/param budget for a 45M model), trained on Colab Pro (a single L4)
across several --resume sessions.
- Init loss 10.9 = ln(50304), the uniform-guess baseline. Landing exactly
there at step 0 is itself a correctness check (
tests/asserts it). - Final loss: train 3.24, val 3.26 (val perplexity ~= 26), tracking each other tightly the whole run: no overfitting, despite ~1.9 passes over the 690M-token corpus.
- Throughput: ~90k tokens/sec steady-state on the L4.
The electric guitar is a direct electric instrument; the other, which tends to be acoustic (not harmonically), while the rest of the material is a direct electric piano. In the concert band, this instrument typically features electric instruments.
Instrumentation and programming The electric guitar has its own distinctive sound that involves creating "the dynamic range" of an ensemble or string section, being used in some music styles such as blues, R&B, soul-rock, black harmony, and electronic music. The technique is also used during recording and recording practice, where it can be performed with a variety of different effects. For example, when the performer performs in front of his audience, the vocalists can mimic the sounds of the instruments by placing them on one's notes (so their playing can produce one's own). When the singer uses fingerings from the solo over the strings, they can use hand pedals such as the pedal point.
The electric guitar had a solid body. The main instrument the amp was on, with the same neck as the other instruments.
Early history (1690-1750) Before that time, the amp used both its lower and upper parts. This era of playing became known by different names in English as "Old England Blues" or "New England Guitar". In 1781, when Captain Richard Balfour died in London, he published his first song called The New American Blues, which is still recorded today. It can be traced back to the original "Piper's Jug", a tune he had collected from the 1660s. From this point on, he experimented with the new, more acoustic guitars over a periodical, which he never played before so long. By then, he was producing and playing for Charles Townshend's Sons Ltd., where he eventually established a publishing house at 35 Southwick Lane.
Explanation: A 45M-param model trained on raw Wikipedia text with a next-token-prediction objective (no fine-tuning, no retrieval, no fact-checking signal) learns the shape of encyclopedic prose extremely well: era headers, specific-sounding dates, named entities, because those are dense statistical patterns in the training corpus. Nothing in the training objective ever rewards a true date over a merely fluent-sounding one, so the model reproduces the form of a Wikipedia article with none of the grounding. That's the "completion model, not a chatbot" limit from 'Honest Expectations' below. The actual fix, already planned in What I'd do differently at scale: fine-tune this frozen base on verified Q&A pairs so "correct" becomes a real training signal, and add retrieval for anything that needs to stay current past the training cutoff.
Errors in Output: The electric guitar wasn't invented until the 1930s, so a header reading "Early history (1690-1750)" is off by roughly two centuries; "Captain Richard Balfour," "The New American Blues," "Charles Townshend's Sons Ltd.," and "black harmony" are fluent, period-plausible inventions that don't exist.
Real failures from building this:
-
Weight tying leaks the current token into its own logits. My init-loss sanity test ("an untrained model should score ~ln(V)") failed at 3.95 vs the expected 4.85. Cause: I'd used
targets=inputs, i.e. asked the model to predict the token it was currently looking at, and with tied embeddings, position t's hidden state starts aswte[token_t], whose dot product with the tied output matrix is largest at…token_t. Even a random model "knows" what token it's standing on. The test now draws targets independently, andtests/test_model.py::test_init_loss_near_uniformdocuments it. -
Vocab padding crashes the tokenizer. I padded the embedding to 50304 for matmul speed;
tiktokenonly knows ids < 50257. An untrained model samples near-uniformly, so the very first checkpoint's sample generation hit a padding id and died withKeyError: 50296, mid-training-run, in the checkpoint code path. Fix:generate(..., vocab_limit=50257)masks the padding ids to −inf before sampling. Lesson: any id the model can emit is an id your decoder must handle, and "can emit" includes rows you never train. -
Wikitext is not one regex. Templates nest (
{{cite {{...}}}}), tables nest, and a single greedy pattern either under-strips (markup soup in the training data) or eats whole paragraphs. The stripper peels innermost constructs in a loop, and I keep a decode-and-eyeball check on the first tokens oftrain.bin, since markup that survives becomes what the model learns.
Other failures from the same silent category: fp16 without
GradScaler underflows gradients (auto-handled here: bf16 where supported,
scaler otherwise), and eval time polluting tokens/sec metrics (timer resets
after eval). Also practical: if this repo lives in OneDrive/Dropbox, exclude
out/ and data/ from sync: a 1GB memmap being re-uploaded every 500 steps
while checkpoints write is misery.
pip install -r requirements.txt
# 1. Build the dataset (streams the dump over HTTP; nothing else touches disk)
python data.py # Simple English Wikipedia, quick start
python data.py --source enwiki --parts 4 # English Wikipedia, the real run (~1B tokens)
# 2. Train (checkpoints, metrics.csv, and sample generations land in out/)
python train.py # ~45M "base" preset
python train.py --model small # ~30M, faster
python train.py --resume # continue from out/ckpt.pt
# 3. Generate
python sample.py --prompt "The electric guitar" --temperature 0.8
# 4. Tests (CPU, <1 min): attention equivalence, causality, RoPE geometry, overfit-one-batch
python -m pytest tests/ -qtrain.py auto-selects CUDA and bf16 when available and falls back to fp16
(pre-Ampere) or CPU/fp32, so the repo runs anywhere, but do the real training
on a GPU. At the default settings one step processes 65,536 tokens; a mid-range
card (RTX 3060-class) does the compute-optimal 0.9B-token budget in about a
day, a 3090/4090 in a few hours.
What those commands do, end to end:
%%{init: {'flowchart': {'nodeSpacing': 16, 'rankSpacing': 25, 'curve': 'linear'}}}%%
flowchart TD
DUMP["Wikipedia Dump<br/>(bz2/XML over HTTP)"] --> STRIP["Markup Stripper<br/>(regex, innermost-out)"]
STRIP --> TOK["GPT-2 BPE Tokenizer<br/>(tiktoken)"]
TOK --> BIN["train.bin / val.bin<br/>(uint16 memmap)"]
BIN --> TRAIN["train.py<br/>(Transformer, AMP, cosine LR)"]
TRAIN --> CKPT["ckpt_best.pt"]
CKPT --> SAMPLE["sample.py"]
Decoder-only Transformer (GPT-style), pre-norm: the same shape as every modern LLM, at 1/10,000th scale.
%%{init: {'flowchart': {'nodeSpacing': 20, 'rankSpacing': 30, 'curve': 'linear'}}}%%
flowchart TD
IN["Input IDs<br/>(B,T)"] --> EMB["Embedding `wte`<br/>(+ dropout) · (B,T,d)"]
EMB --> X["x"]
subgraph BLOCK["Transformer Block ×6"]
direction TB
X --> LN1["LayerNorm"] --> ATTN["Multi-Head Causal Self-Attention<br/>RoPE on Q,K · + output proj"]
ATTN --> ADD1((+))
X --> ADD1
ADD1 --> LN2["LayerNorm"] --> FFN["Feed-Forward<br/>Linear → GELU → Linear (4×)"]
FFN --> ADD2((+))
ADD1 --> ADD2
end
ADD2 --> LNF["LayerNorm (ln_f)"]
LNF --> HEAD["LM Head (tied) · (B,T,d)"]
HEAD --> LOGITS["Logits<br/>(B,T,50304)"]
style BLOCK fill:#dbe9ff,stroke:#4a6fa5,stroke-width:1px,color:#14213d
| Choice | What I did | Why |
|---|---|---|
| Norm placement | Pre-norm (norm inside the residual branch) | Post-norm (the original 2017 layout) makes the residual stream pass through LayerNorm, which destabilizes deeper stacks; pre-norm keeps an identity path and trains without fragile warmup tuning. |
| Positions | RoPE (rotary embeddings), implemented by hand | Rotating Q/K channel pairs by position makes attention scores depend only on relative offset, no learned position table needed, and it's what Llama-class models use. tests/ verifies both the rotation property and the relative-position property. |
| Attention | Manual softmax(QKᵀ/√d + mask)V, with an optional fused-kernel path |
The hand-written path is the reference; flash=True dispatches to F.scaled_dot_product_attention for the actual training run (same math, fused kernel; a test asserts they agree to 1e-5). Softmax runs in fp32 even under bf16 autocast. |
| LayerNorm | Written out (mean/var/rsqrt), stats in fp32 | The variance of bf16 activations is where mixed-precision LMs quietly rot. |
| Embeddings | Tied input/output matrix | At this scale the 50304×512 embedding is 58% of all parameters, and tying saves 25M params and regularizes. (It also bit me; see What broke.) |
| Vocab | GPT-2's 50257, padded to 50304 (×64) | GPU matmuls are measurably faster when the vocab dim is a multiple of 64. (This also bit me; see What broke.) |
| Init | N(0, 0.02), residual projections scaled by 1/√(2·n_layers) | Each block adds two branches to the residual stream; scaling their output init keeps activation variance ~1 at depth instead of growing linearly. |
| Optimizer | AdamW (β₂=0.95, wd=0.1 on matrices only), cosine LR + warmup, grad clip 1.0 | GPT-3 settings, standard at this scale. Weight decay is not applied to LayerNorm gains/biases, since decaying a normalization scale toward zero fights the normalization. |
| Data loading | np.memmap over a flat uint16 token stream, random windows |
No DataLoader, no padding, no epochs: every offset in a continuous token stream is a valid training example, and the OS page cache does the buffering. uint16 halves the disk footprint (50257 < 2¹⁶). |
data.py streams pages-articles dumps straight from dumps.wikimedia.org over
HTTP → bz2 → XML, strips wiki markup with a hand-rolled set of regexes
(templates and tables are nested, so they're peeled innermost-out), and
tokenizes into a flat binary. Articles whose categories or infoboxes look
music-related (a keyword classifier: album, song, composer, jazz, guitar,
symphon-, …) are written 3× to the training stream. Since training samples
random windows, duplication is upweighting, so no sampler changes are needed.
Validation articles are held out entirely and never duplicated, so val loss
stays honest.
Honest expectations: a 45M-param base model is a completion model, not a chatbot. Trained on this mix it should continue "The saxophone is" with coherent, music-literate encyclopedic prose, but it will not reliably answer "who wrote Kind of Blue?". Getting question-answering behavior would take a small supervised fine-tune on Q&A pairs afterward (see below), which the same training loop could do with a different data file.
- Train my own tokenizer. GPT-2's BPE is the one off-the-shelf component here, and it's frozen to 2019 web text. At scale, vocab is a tuning knob (size, byte-fallback, domain coverage), and I'd train BPE on my own corpus.
- KV cache for generation.
generate()recomputes the full forward pass per token, O(T²) total. Fine for 200-token samples during training; the first thing to fix for real inference. - Real data pipeline. Regex-stripping wikitext is "good enough for pretraining", not good: I'd use a proper parser, dedupe near-duplicates (Wikipedia is full of formulaic pages), and mix sources rather than oversampling one domain by keyword matching. Data quality is the highest-leverage knob at every scale.
- Distributed training. This repo is deliberately single-GPU. The next step is DDP (data parallelism), then, past ~1B params, sharded optimizers (ZeRO/FSDP), since the optimizer states are 2× the model in fp32.
- Proper eval. Val loss and eyeballing samples is fine at 45M. At scale: held-out perplexity per domain, and task benchmarks, tracked per checkpoint.
- A fine-tuning stage. To actually answer music questions: freeze this
as the base model, build a few thousand Q&A pairs, and run the same loop at
low LR on formatted examples. The infrastructure here already supports it,
it's just a different
train.bin.
- Vaswani et al., Attention Is All You Need (2017): the transformer
- Radford et al., Language Models are Unsupervised Multitask Learners (2019): GPT-2, decoder-only recipe, init scaling
- Su et al., RoFormer (2021): rotary position embeddings
- Hoffmann et al., Training Compute-Optimal LLMs (2022): the ~20 tokens/param rule used to size the token budget
- Karpathy's nanoGPT: the spiritual ancestor of every small-GPT repo, including this one; the memmap batching trick is his