10h 27m 16s logged

modern-lm: modernizing the architecture (RMSNorm, RoPE, SwiGLU)

Starting from my GPT-2 124M reproduction and dragging the architecture from 2019 to 2026, one component at a time. The plan is a ladder: swap one thing, understand exactly what it does, keep going. This devlog covers the three architecture swaps.

What changed:

LayerNorm to RMSNorm, and dropped every bias. RMSNorm skips the mean-subtraction and only rescales by root-mean-square times a learned gain. The residual stream’s problem is magnitude drift, not offset, so re-centering was doing nothing. Half the ops, one fewer param vector per norm, and it sets up Muon later (which wants a model that is purely matrices plus a few gains).
Learned position embeddings to RoPE. Deleted the wpe table entirely. Instead of adding a position vector, RoPE rotates the query and key feature pairs by a position-dependent angle, so attention scores end up depending on the relative distance between tokens, not their absolute index.
GELU MLP to SwiGLU. Replaced the Linear to GELU to Linear block with a gated FFN: one branch computes features, a second branch (through SiLU) decides how much of each passes, then they multiply and project down. Went with 4x width, so it’s a bigger model, not a param-matched swap.

Deliberate: I care more about a stronger model here than a clean ablation, and I can spend the compute.

What was actually hard:

The swaps are each ~20 to 40 lines. The trap was RoPE. It’s the classic silent bug: get the pairing convention wrong between the cos/sin cache and the apply function and the model still trains fine, it just quietly underperforms and you never find out. So I wrote a unit test instead of trusting the loss curve: shift the whole sequence by k positions and assert the attention scores between the same token pairs are unchanged. Rotation invariance either holds or it doesn’t (like an odometer).

GPT-2 124M Reproduction

Reproduced GPT-2 124M (the old 2019 model) from scratch, trained on ~10B tokens across 2x RTX 5090 in ~7 hours. Final val loss ~3.10.

Dataset

FineWeb-Edu sample-10BT, ~10B tokens
GPT-2 BPE (tiktoken), vocab 50257 padded to 50304
100 uint16 shards (99 train, 1 val), ~19GB
Streaming download + multiprocessing tokenizer — switching to imap_unordered with a bigger chunksize took tokenizing from ~2 hours to ~6 minutes

Architecture

GPT-2 124M: embd=768, block_size=1024, weight-tied embeddings (~124M params)

12 blocks, each with causal self-attention (12 heads, flash attention) + MLP (Linear → GELU → Linear, 4x)
Pre-LayerNorm + residuals on both sublayers

Training

bf16 + torch.compile, AdamW, OneCycleLR cosine with warmup, grad clip 1.0
0.5M-token effective batch (524,288), held fixed via gradient accumulation
DDP across 2x RTX 5090 on RunPod, 19,073 steps (~1 epoch), ~7 hours
Checkpoint every 200 steps + auto-resume babysit script

Results

Val loss ~3.10
Coherent, topical English continuations. It’s a base model, so it continues a prompt rather than answering it .
Shipped an inference script + Gradio demo, uploaded to Hugging Face

What was actually hard

The architecture was basically the char transformer from last devlog, just wider (124M vs ~10M) with BPE instead of 65 chars. The model was the easy part.

The real work was scale engineering: sharding 10B tokens to stream off disk, gradient accumulation to hit a 0.5M-token batch on cards that OOM at batch 8, multi-GPU gradient sync with DDP, mixed precision, and checkpoint/resume so a 7-hour cloud run survives interruptions. That’s what turns a transformer into a trained LM.

Open comments for this post

@xerneas on Makemore · about 1 month ago

11h 7m 7s logged

Tiny Shakespeare (decoder-only transformer)

Dataset:

Tiny Shakespeare, ~1.1M characters
vocab 65

Architecture: GPT-style decoder

token embedding + learned position embedding (embd=384)
6 transformer blocks, each:
- multi-head causal self-attention (8 heads, head_size=48)
- feed-forward (Linear → ReLU → Linear)
- pre-LayerNorm + residual connections on both sublayers
final LayerNorm → linear to vocab
block_size=64, dropout=0.2

Training:

AdamW, lr=1e-4
ExponentialLR, gamma=0.99995
batch size 1024
5000 iters
CUDA + torch.compile

Results:

train loss = 1.26
val loss = 1.52

Why it worked:

Attention replaces flattening. Instead of mashing all 64 embeddings into one vector, each position attends to the relevant earlier positions directly. Sequence structure is preserved instead of thrown away — the exact failure mode of mlp3.
The receptive field is the full block (64 chars) and every position can use all of it, not just a fixed local window.
Residual connections + pre-LayerNorm let 6 blocks stack and actually train.
Position embeddings give the model order information that a bag-of-embeddings MLP never had.

Against the mlp3 goals:

Goal was dev loss < 1.7 → hit 1.52.
Goal was “generate actual words and sentences” → it now produces speaker tags, line breaks, and mostly real English words in Shakespearean cadence. It learned the SPEAKER:\n dialogue form like mlp3 did, but this time it fills the dialogue with coherent structure instead of gibberish.

The train/dev gap:

1.26 vs 1.52→ much tighter than mlp3’s 1.12/2.02. Dropout 0.2 + a context long enough to actually be useful means it’s learning real patterns, not memorizing short-range ones.

Notes / loose ends:

LR never decays much → gamma=0.99995 over 5000 steps only drops lr from 1e-4 to ~7.8e-5, so the schedule is nearly flat. Loss was still falling at step 5000; more iters or a higher/decaying lr would likely push val below 1.5.

Makemore