gpt2-muon-124m

A faithful GPT-2 (124M) trained from scratch on 10B tokens of FineWeb-Edu, using a Muon-split optimizer + warmup-stable-decay (WSD) recipe instead of the usual AdamW + cosine. The architecture is byte-identical to OpenAI GPT-2 small — only the training recipe differs.

Part of SkyAI, a from-scratch LM project where every layer, optimizer step, and tokenizer decision is hand-written.

Results

metric	gpt2-muon (this)	GPT-2 124M	nanoGPT / llm.c 124M (AdamW)
val_loss (FineWeb-Edu val shard)	2.9653	~3.29	3.28
HellaSwag (acc_norm)	0.3238	0.294	~0.30
HellaSwag (acc)	0.2991	—	—
LAMBADA (acc / ppl)	0.2826 / 27.81	~0.33 / ~35	—

Protocol note. val_loss is 2.9653 scored on the full 100M-token FineWeb-Edu validation shard. An earlier version of this card reported 2.99, which was the in-training figure computed on a 10.5M-token prefix — the harness counted validation micro-batches rather than tokens, so the budget depended on batch and world size. Absolute val_loss is protocol-sensitive; always compare on a matched window.

Evaluated with Karpathy's exact HellaSwag method (10,042 val examples), on the same data, val shard, and tokenizer as the references. The Muon+WSD recipe clearly beats the AdamW+cosine reference on val_loss and HellaSwag at matched data and scale.

Compared head-to-head against a modern architecture

This model is rung 2 of a three-rung ladder. Rung 3, muteptr/skyai-modern-xs, keeps the tokenizer, data, token budget, and schedule shape identical and changes only the architecture (RMSNorm, RoPE, SwiGLU, GQA, QK-norm, untied embeddings, logit soft-cap). On a matched 100M-token validation shard it reaches 2.9548 vs 2.9653 here (+0.0104 nats, paired 95% CI [+0.0102, +0.0107]) and LAMBADA ppl 26.25 vs 27.81, at ~8% fewer FLOPs/token — while being statistically tied on all accuracy benchmarks.

Caveats (read these)

• Recipe-level, not Muon in isolation — Muon, WSD, and tuned LRs/weight-decay moved together; the WSD decay-to-zero contributes a real share of the val-loss margin.
• Single seed.
• val_loss vs GPT-2 is home-field — this model trained on FineWeb-Edu and GPT-2 did not, so the neutral comparison is HellaSwag.
• Base model, not instruction-tuned, and trained on educational text — strongest on encyclopedic/explanatory prompts.

Usage

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("muteptr/gpt2-muon-124m")
model = AutoModelForCausalLM.from_pretrained("muteptr/gpt2-muon-124m").eval()

inputs = tok("Photosynthesis is the process by which", return_tensors="pt")
out = model.generate(**inputs, max_new_tokens=40, do_sample=True, top_k=50, temperature=0.8)
print(tok.decode(out[0], skip_special_tokens=True))

Sample output:

Photosynthesis is the process by which a plant absorbs carbon dioxide from the air, stores it in the soil, and uses the stored carbon to make its own food. The main function of Photosynthesis in a plant is to capture sun…

Training

Data	FineWeb-Edu sample-10BT (~10B tokens, gpt2 BPE) — 99 train shards + 1 val
Hardware	8×A100-80GB SXM4, ~2.1h
Optimizer	Muon-split — Newton-Schulz orthogonalized momentum on 2D hidden matrices, AdamW on embeddings / norms / biases. Re-tuned for gpt2 (tied wte + LayerNorm): embedding_lr=0.006, matrix_lr=0.015, weight_decay=0.28
Schedule	warmup-stable-decay — 715 warmup / 19,073 steps, LR decayed to 0 over the final 40%
Batch	524,288 tokens/step (0.5M), context length 1024
Precision	bf16 autocast (fp32 master weights; released in fp32)

Full, runnable recipe: configs/gpt2-muon.yaml.

License

MIT for the weights. Training data is FineWeb-Edu (ODC-BY).