Modern LM

Modern LM is a 481M parameter language model I built from scratch, one modern component at a time. It picks up exactly where my Zero to GPT-2 project left off: that one reproduced a 2019-era GPT-2 by hand, and this one asks what actually separates a modern model from GPT-2, then implements each of those ideas from scratch and trains the whole thing end to end. I trained the base on 10 billion tokens and then post-trained it into a chat assistant with calculator tool-use, closed-book QA, and retrieval-grounded answering.

Why I built it

Zero to GPT-2 ended with a working but dated model, and I said the natural next step was to apply the fundamentals to current architectures rather than keep polishing a 2019 design. This is that step. I read through a stack of papers, tech reports, and speedrun writeups on what makes a modern model, picked the parts that looked most promising, implemented each one myself without copying a reference, checked that it actually helped, and got all of them to cooperate in one clean codebase. The full lineage back to Zero to GPT-2 is in the acknowledgments of the repo.

None of the ideas are mine. The work was reading each one, building it by hand, verifying it, and stitching them together. It is a small model, so the format and behavior are real but the world knowledge is limited. That tradeoff is the whole point of building it yourself.

What I picked, and why

Each of these is a modern component I pulled in and implemented from scratch. The reason each earns its place is next to it, with the paper I worked from:

• RMSNorm with no biases: normalizes by root-mean-square only and drops every bias term. I picked it because it matches LayerNorm's stability at lower cost, and dropping biases is a small, free simplification that modern models have converged on.
• RoPE : rotary position embeddings rotate query and key pairs by position. I chose it over learned absolute positions because it encodes relative distance directly and generalizes better to sequence lengths past what it trained on.
• SwiGLU feed-forward: a gated FFN with three matrices. I used it because the gating consistently buys a quality improvement over a plain two-matrix MLP for the same parameter budget.
• QK-norm , used at scale in ViT-22B : per-head RMSNorm on queries and keys. I added it to keep attention logits from growing and destabilizing training, which lets me train at a higher learning rate.
• GQA : 16 query heads share 4 key-value heads. I picked it specifically for inference: it shrinks the key-value cache roughly four times with little quality cost, which matters far more than the parameter count.
• The Muon optimizer on the 2D hidden matrices: it orthogonalizes the momentum-averaged gradient. I chose it because in the modded-nanogpt speedruns it reaches a target loss in noticeably fewer steps than AdamW, and it moved my loss more than any other single change.
• A WSD learning-rate schedule (warmup, stable, decay): I picked it over cosine because the long stable phase can be extended without re-planning the schedule, and the final decay to zero delivers the single largest drop in loss, which the curve below shows clearly.
• FlashAttention (v2 , via PyTorch SDPA): fused, memory-efficient attention. I used it so the 2048-token context fits and runs fast without a custom kernel.
• Untied input and output embeddings with zero-init residual writers (tying background , zero-init from modded-nanogpt ), plus value embeddings, U-net style skip connections between mirrored layers (from the speedrun Field Guide ), and a logit softcap of 15 tanh(logits/15) from the Gemma 2 report . These are small, well-attested wins I stacked on top.

Why the optimizer is split

Muon only makes sense on the 2D hidden weight matrices, where orthogonalizing the update is meaningful. So the optimizer is split: Muon runs on those matrices, and AdamW handles the embeddings, the LM head, and the one-dimensional norm and scalar parameters. The embeddings are excluded from Muon by identity. Getting that split right was one of the fiddlier parts of making all of these cooperate.

Architecture

• Parameters: 480.9M, 24 layers, model dimension 1024.
• Attention: 16 query heads, 4 key-value heads (GQA), 2048 context length.
• Vocabulary: 50,304 (GPT-2 byte pair encoding, padded), with reserved special tokens for chat, thinking, and tool calls used only in post-training.
• Precision: bf16 autocast throughout.

What I started with, and where I ended up

Like the GPT-2 run, this starts at a loss around 11.0, guessing uniformly over the vocabulary. Over the full run the validation loss falls to 2.7618. The most interesting part is the schedule: the WSD decay phase at the very end, where the learning rate is annealed to zero, delivers the single largest drop in loss.

The training run

• Data: FineWeb-Edu 10B sample (82%) mixed with clean Python from CodeParrot (18%), about 10 billion tokens total. I added code so the model would pick up some programming ability, not just prose.
• Schedule: WSD, with 3% warmup, a long stable phase, then a linear decay to zero. Muon at 0.02 and AdamW at 1.2e-3.
• Batch: a fixed 524,288 tokens per optimizer step via gradient accumulation, so the dynamics stay constant regardless of hardware.
• Hardware: a single H100 80GB at roughly 131,000 tokens per second, about 21 hours end to end.

Post-training: the fine-tunes

The base model only predicts the next token, it is not an assistant. Every model below comes from the same pretraining run, then fine-tuned. Fine-tuning here means continuing to train the exact same 481M network on new, task-specific data with a masked objective, so I am not adding parameters, just teaching the existing weights a new behavior. The lineage is one chat model branched into three specialists. For the recipe I leaned on nanochat and TRL .

1. Supervised fine-tuning: the chat model

This is the foundation the others branch from. I took the base model and fine-tuned it on conversations formatted as ChatML, using the special tokens I reserved in the vocabulary during pretraining. The key detail is masking: the loss is applied only to the assistant's tokens, not the user's prompt. That is what teaches the model to answer and to stop at the end of a turn, rather than continue or echo the prompt back. I also mixed in synthetic arithmetic tool-call examples at this stage, so the chat model has light tool use before I even specialize it. The result is modern-1024x24-sft .

2. Calculator tool-calling

Branched from the chat model, this one is a short, focused fine-tune on writing tool calls. The model learns to emit a call, a real calculator evaluates it, and the result is fed back for the model to use in its answer. What I did to make this honest: during training the injected tool result is masked out of the loss, so the model is rewarded for writing the correct call and using whatever comes back, not for memorizing the answer. The calculator itself is a strict arithmetic evaluator that rejects names, function calls, and imports, so it is a calculator and never code execution. The result is modern-1024x24-sft-toolcall .

3. Closed-book question answering

Also branched from chat, this is a short fine-tune on question and answer pairs where the model must answer from what is baked into its own weights, with no retrieval or tools. It is the honest test of what the base actually learned during pretraining. At 481M parameters and 10B tokens the format is reliable, the model answers cleanly and concisely, but the facts are limited by how much a small model can store. The result is modern-1024x24-sft-qa .

4. Retrieval-grounded QA (RAG)

The last branch, and the most interesting. Here the model is given a short context plus a question and must answer only from that context. What I did specifically: I trained it on a mix that includes questions whose answer is not in the context, so it learns to say I don't know and abstain rather than hallucinate. That abstention is the whole point of a grounded model. It also produced my favorite lesson of the project: this model once reached the lowest loss of any stage, purely by learning to always say I don't know, which is a sharp reminder that a low loss is not proof a model works, and why I test with questions whose answers I already know. The result is modern-1024x24-sft-rag .

What I learned

• Most of these ideas are small, independent wins that compound. Muon and the WSD decay phase moved the loss the most.
• Most of a model's usefulness is unlocked in post-training, but post-training cannot add knowledge the base does not have. At this size the format is reliable and the facts are not.
• A low loss is not proof a model works, as the RAG abstention result showed.
• Implementing each idea by hand and then diffing against a reference taught me far more than reading the papers would have.

What I left out, and why

Reading widely also meant deciding what not to include. A few things I studied and deliberately left out of this model:

• Multi-head Latent Attention (MLA) from DeepSeek-V2: it compresses the key-value cache more aggressively than GQA. I left it out because GQA already gave me the inference-memory win I wanted, and MLA adds real implementation complexity for a payoff that mostly shows up at much larger scale and context than this model.
• MuonClip and QK-clip from the Kimi K2 tech report: these target training instability that appears at very large scale. At 481M with QK-norm already in place, my training was stable, so the extra machinery would have been solving a problem I did not have.
• Engram: interesting but niche, and outside the scope of a first modern build. I wanted a clean, well-understood stack, not every idea at once.

The theme is the same across all three: I kept the model focused on components whose benefit I could actually see at this size, rather than adding tricks aimed at models orders of magnitude bigger.

What's next

The obvious steps I have not done yet are preference tuning with DPO to make the assistant's answers more aligned with what people actually prefer, and reinforcement learning with verifiable rewards in the spirit of open-r1 for tasks with checkable answers like math. Both build directly on the post-training stack that is already here.

Try it

All five models (base, chat, calculator, closed-book QA, and RAG) are published on Hugging Face , and a Streamlit playground wraps them behind one interface. The full architecture, training, and post-training code is in the repository .