ulmentflam/gpt2-774m-fineweb-fp8-mojo

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

A GPT-2 774M (36 layers, 20 heads, 1280 channels (d36)) 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

Training

• —Precision: fp8 (e4m3) for every per-block linear GEMM, with delayed scaling from a 16-step amax history; parameters stay bf16 and the optimizer stays fp32.
• —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 2.5e-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-774m-fineweb-fp8-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.