gpt2-muon-124m

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

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

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

License

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