Reproducing OLMo 3 7B Pre-training in MaxText

OLMo 3, developed by the Allen Institute for AI (AI2), is a state-of-the-art, fully open language model trained with a modern architecture and a multi-stage training recipe. To evaluate the capabilities of MaxText on Google Cloud TPUs, our team set out to reproduce AI2’s OLMo 3 7B from scratch. We chose OLMo 3 because it combines three properties that rarely appear together. It is a strong, modern 7B model trained at real production scale. AI2 exposes nearly the complete model flow, including data, code, configurations, checkpoints, logs, and evaluations. And finally, it gives us an independent PyTorch and GPU reference against which we can test MaxText and TPUs.

PyTorch → JAX model conversion. OLMo 3's architecture (reordered-norm block, QK-norm, 3:1 sliding/global attention) ported to MaxText and verified with a logit-parity check: the converted step-0 checkpoint matches the HuggingFace reference at KL ≈ 1.5e-3, the "same model, different framework" noise floor, and at the full 8192-token context in bfloat16 the two agree on the top-1 token 98.75% of the time.

Verification that catches real bugs. Held-out evals caught a data-loader bug that made MaxText look like it was beating the reference; the gain was memorization.

Reliability over a multi-week run. Checkpoint-and-resume replays the run exactly: a controlled A/B shows Δ = 0.000 at every step after a resume, and when a host failure killed the stage-2 run mid-flight, the resumed run re-trained 127 steps at Δ = 0.000 in logged loss and perplexity.

Resizing the training job mid-flight. At step ~1.05M we lost three quarters of our capacity; the run resumed on a slice one quarter the size with no recipe change (same script run_olmo3_7b_stage1.sh scales per device batch size to keep GBS constant), per-device throughput preserved within 1% (≈100% strong scaling, measured in both directions).

Changing TPU generation mid-recipe. Stage 2 pointed the identical launcher at v5p instead of Ironwood, changing only the device type, and sustained 57.4% MFU.

Performance work that paid for itself. 44.5% MFU on Ironwood at 7B via SparseCore collective offload, remat tuning, and optimal sharding, roughly a third of the compute budget bought back.

Co-design for TPU: faster at the same quality. Reshaping attention from 32 heads × head-dim 128 to 16 × 256, at identical parameters and FLOPs, runs +12.4% faster because head-dim 256 fully utilizes Ironwood's 256×256 MXU, and its loss curve matches the original through 120B tokens (30k steps). This was a side ablation; the reproduction kept the original architecture.