GPT-2 124M (nvfp4), 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.3135
HellaSwag (acc_norm) 2982/10042 = 29.70% (Wilson 95% CI [28.8%, 30.6%])
Final train loss 3.3445
Throughput ~649,137 tok/s (706.7 ms/step, median over the last 500 steps)

Training

  • Precision: NVFP4 (e2m1, block-scaled) on the middle blocks' MLP forward GEMMs with stochastic rounding and a randomized Hadamard transform; backward stays bf16, parameters stay bf16, 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 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-nvfp4-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.

Downloads last month
24
Safetensors
Model size
0.1B params
Tensor type
BF16
Β·
Inference Providers NEW
This model isn't deployed by any Inference Provider. πŸ™‹ Ask for provider support

Dataset used to train ulmentflam/gpt2-124m-fineweb-nvfp4-mojo

Collection including ulmentflam/gpt2-124m-fineweb-nvfp4-mojo