A 51M-parameter GPT,
trained from scratch.

The model, end to end.

A decoder-only transformer — nanoGPT modernized with ideas from Gemma: a token embedding, eight pre-norm blocks, a final norm, and an output projection that reuses the embedding.

What each piece does.

• RMSNormkeeps numbers in scale
  Stops the values flowing through the network from growing or shrinking out of control, which keeps training stable. It's a simpler LayerNorm — it skips the mean-centering and bias, so fewer parameters and a bit faster.
• RoPEgives it word order
  Attention on its own ignores order. RoPE fixes that by rotating each word's vectors by an angle based on its position — so the model can tell how far apart two words are. No extra parameters, and it handles text longer than it was trained on.
• QK-normkeeps attention from exploding
  Attention scores come from multiplying two vectors together; if those vectors get too large the scores blow up and training destabilizes. QK-norm rescales them first so the scores stay in a healthy range — cheap insurance for a from-scratch run.
• SwiGLUprocesses each word
  The layer that transforms each word on its own. It's "gated": one path decides how much of another passes through, letting the model learn which features matter — more capable per parameter than a plain layer. down(silu(gate·x) * up·x).
• Weight tyingreuse one table
  Turning words into vectors (input) and vectors back into word-scores (output) both need a big lookup table. Sharing one table for both — instead of two — cuts the model roughly in half and tends to improve quality.

From raw text to a chat model.

Two phases — pretrain, then fine-tune — on the same machinery; only the data and objective change. The shared choices first, then each phase.

Shared foundations

• bf16 mixed precision
  Math runs in bfloat16 — about 2× faster and half the memory of fp32 — while a full-precision master copy of the weights keeps the updates stable.
• Gradient accumulation
  Sum gradients over several micro-batches before each weight update, so a single A100 reaches a large effective batch without running out of memory.
• AdamW
  β 0.9 / 0.95, weight decay on matmul weights only, gradient clip 1.0, cosine LR with warmup. Same optimizer family every phase, retuned per phase.
• Checkpoint & resume
  Checkpoints to Google Drive store weights + optimizer + step + data cursor, so a dropped Colab session continues from exactly where it stopped.

1 · Pretrain

done

• Corpus
  70% FineWeb-Edu · 20% Cosmopedia · 10% FineMath, streamed and packed into one uint16 stream with no padding.
• Budget
  3.5B tokens ≈ 13.3k steps at 256k tokens/step (~one epoch). That's ~136 tokens/param — roughly 7× past Chinchilla-optimal, the over-training strategy small models like SmolLM use. ~3h on one A100.
• Settings
  Context 1024, cosine LR 6e-4 → 6e-5, 2k-step warmup. Objective: next-token prediction over the whole sequence.

2 · SFT

done

• Corpus
  Science Q&A — SciQ + OpenBookQA + ARC (~22k) — plus a small smol-smoltalk chat subset (~10k), all rendered to one question → answer shape in a chat template.
• Template
  ChatML markers as plain text; each assistant turn ends with <|endoftext|> — the token the base already knows as a stop — so the vocabulary stays fixed: no new tokens, no checkpoint surgery.
• Budget
  A few epochs over the small set — a few hundred optimizer steps, not billions of tokens. Minutes on one GPU.
• Settings
  Initialized from the base weights with a fresh optimizer, cosine LR ~5e-5 (far gentler than pretrain), loss on assistant tokens only — so it learns to answer, not echo the question. The lowest-val checkpoint is kept.

How we know it learned.

Scored on the SciQ test set: science questions, each with one correct answer and three distractors. The chat model only — the base can't do the task.

• Multiple-choice by likelihoodnot string-matching
  Letting a 51M model write an answer and matching the text is too brittle ("carbon dioxide gas" vs "carbon dioxide"). Instead we ask which of the 4 options the model finds most probable as the answer — one forward pass each — and count it right when the true one wins. It's the inference-time mirror of the masked SFT loss.
• Length-normalized accuracythe headline
  Score each option by its average log-probability per token, so the model can't just always prefer the shortest answer. A raw-sum variant rides along to flag when answer length is driving the picks.
• 25% random baselinethe floor
  Four options means guessing scores 25%. That's the bar every result is measured against — did we beat chance?

Result

Length-normalized accuracy on the held-out split (raw-sum agrees at 47.8%, so length isn't driving it) — nearly twice chance for a 51M model from scratch. Free-form generation is rougher: it lands the right topic, not always the exact word.

How it picks each word.

Generation is autoregressive: feed the context, get a probability for every token, pick one, append, repeat. These knobs shape that pick — and a KV cache makes the repeating fast.

• Temperatureflat vs spiky
  Scales the logits before softmax. Low (<1) sharpens the distribution toward the likeliest tokens — safe but repetitive; high flattens it — more varied but riskier. Set to 0 it becomes greedy: always take the single top token.
• Top-pnucleus sampling
  Keep only the smallest set of most-likely tokens whose probabilities add up to p (e.g. 0.9), then sample from those. Cuts off the unlikely long tail that produces gibberish, while staying flexible when the model is genuinely uncertain.
• Repetition penaltybreaks loops
  Small models love to loop. Tokens already in the context get their scores pushed down before sampling, so the model is discouraged from repeating itself. Default 1.2 (1.0 = off).
• KV cachedon't redo old work
  A token's attention needs the keys and values of every earlier token — but those never change once computed. We store them in a fixed-size buffer and reuse them, so each step processes just the one new token instead of re-running the whole context. A JIT-compiled single-token step keeps the array shapes constant, so it compiles once and stays fast. On by default; --no-cache falls back to the full re-run.
• Sliding windowgenerate past the context
  In the --no-cache path, only the last 1024 tokens are fed each step, so generation can run longer than the context window — the oldest tokens scroll out as new ones arrive. (With the fixed-size cache on, context is instead capped at the window.)

What 51M can and can't do.

It clears the bar it was built for — 47% on SciQ vs 25% chance. Using it also makes the capacity ceiling concrete, and it's worth being honest about where that line falls.

• What it learnedthe win
  Fluent English, the chat format, and a real slice of grade-school science. Ask the right kind of question and it answers — "what galaxy is our solar system part of?" → the Milky Way; "dialysis treats failure of what organs?" → kidneys. That's genuine knowledge, surfaced by fine-tuning, on a model trained from nothing.
• Knowledge is thinthe ceiling
  Parameters are where facts live, and 51M holds few. Outside its trained slice it confabulates confidently; even on-topic it often grabs a neighbor — "opposite of melting?" → a melting point, "potential difference in a circuit?" → current (not voltage). It knows the area, not always the exact fact.
• Understanding > vocabularya nice surprise
  Sometimes it grasps the concept but reaches for the wrong word — "sites of protein synthesis?" → protein factories, the textbook nickname for ribosomes. It learned the ideas more robustly than the precise terms — which is exactly why we score by ranking the options, not exact-match generation.
• Fixable vs needs scalethe takeaway
  Decoding quirks (looping, echoing the question) are fixable with sampling settings. The real limits — hallucination, open-ended conversation, breadth across domains — aren't a recipe problem, they're capacity. SFT surfaces what pretraining stored; it can't add what isn't there. Those gains come from a bigger model and more pretraining: the path to a v2.

Built on prior work.