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)	2.99	~3.29	3.28
HellaSwag (acc_norm)	0.324	0.294	~0.30
HellaSwag (acc)	0.299	—	—
LAMBADA (acc / ppl)	0.283 / 27.8	~0.33 / ~35	—

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.

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("MuteBuster/gpt2-muon-124m")
model = AutoModelForCausalLM.from_pretrained("MuteBuster/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).