ulmentflam/gpt2-124m-fineweb-mojo

GPT-2 124M (bf16), trained from scratch on FineWeb with llm.mojo

A GPT-2 124M (12 layers, 12 heads, 768 channels (d12)) language model trained from scratch — random initialisation, no GPT-2 warm-start — for one epoch of the FineWeb classic 10B-token sample, using llm.mojo, a Mojo/MAX port of Andrej Karpathy's llm.c.

This is a training-run artifact and reproducibility reference, not a state-of-the-art model. At this scale and data budget GPT-2 produces locally-coherent but not highly capable text.

Why these weights were replaced

An earlier version of this repo held a checkpoint trained with a defect: a scratch-buffer overrun in the fused bias-gradient kernel meant several bias tensors received no gradient at any step. It was measurable directly from the published weights — GPT-2 initialises biases to exactly zero, so a bias still bit-exactly zero after 22,345 optimizer steps was never updated.

The kernel is fixed and this checkpoint is a complete re-run. The export now refuses to publish any checkpoint whose matmul biases are all zero. Full write-up: docs/ai/dbias_scratch_overrun_silent_zero_bug.md.

The previous weights remain reachable in this repo's commit history.

Results

Final validation loss	3.2869
HellaSwag (acc_norm)	3010/10042 = 29.97% (Wilson 95% CI [29.1%, 30.9%])
Final train loss	3.3182
Throughput	~771,577 tok/s (594.5 ms/step, median over the last 500 steps)

Training

• Precision: bf16 mixed precision — parameters, activations and gradients in bf16, fp32 AdamW master weights and moments.
• Data: FineWeb classic sample-10BT, GPT-2 BPE, one epoch = 22,345 steps
• Batch: 458,752 tokens/step (7 ranks x gradient accumulation)
• LR: cosine, peak 6e-4, 700-step warmup, decayed to 0, weight decay 0.1
• Parallelism: 7x RTX PRO 6000 Blackwell Max-Q, ZeRO-1 data parallel
• Code: github.com/ulmentflam/llm.mojo

Files

• model.safetensors — GPT2LMHeadModel-compatible export, loadable with transformers
• model_22345.bin — the raw llm.mojo/llm.c-format checkpoint

infer_gpt2.mojo loads any of a local .bin, a local .safetensors, or this repo directly via --hf ulmentflam/gpt2-124m-fineweb-mojo.

Reproducing

git clone --recurse-submodules https://github.com/ulmentflam/llm.mojo.git
cd llm.mojo && make install-cuda
pixi run python data/fineweb.py -t classic -v 10B -m gpt-2
make build-bf16 WORLD_SIZE=7

See the launcher for this arm's exact flags.