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).

Next

• QK-norm, untied embeddings + zero-init prWSD schedule.
• Then the actual point: benchmark the modernized model against the untouched GPT-2 baseline (Row 0) at a fixed
  token budget on FineWeb-Edu, and see which
• Then scale the frozen recipe to ~500M.