Google reproduces AI2's OLMo 3 model on TPUs, proving JAX can match PyTorch at scale

Researchers at Google have successfully reproduced the OLMo 3 7B language model—developed by the Allen Institute for AI—using MaxText, a JAX-based training framework designed for TPUs. This feat demonstrates that high-performance language model training can be faithfully transferred from PyTorch on GPUs to JAX on Google Cloud TPUs, and that the match holds not just on training loss but on downstream metrics that matter.

Why OLMo 3 Matters as a Test Case#

The team chose OLMo 3 because it combines three rarely-seen properties. First, it is a capable, production-scale 7-billion-parameter model with modern architecture (including non-standard choices like reordered-norm blocks, QK-normalization, and mixed sliding-window/global attention). Second, AI2 has released nearly complete transparency—data, code, configurations, checkpoints, and published evaluation results via Weights & Biases. Third, it provided an independent PyTorch-on-GPU reference that could serve as ground truth. The core question: can a recipe trained on one framework-hardware combination be reproduced faithfully on another, and how do you actually prove it?

Technical Achievements: From Model to Data Pipeline#

The reproduction required building OLMo 3 inside MaxText from scratch. The team:

• Ported the architecture (reordered-norm block, QK-norm, 3:1 sliding/global attention pattern) and verified it with a logit-parity check against HuggingFace weights—the converted checkpoint agreed with the reference at KL ≈ 1.5e-3, the expected "same model, different framework" noise floor.
• Implemented a skip-step optimizer matching OLMo-core's exact semantics, including Bessel-corrected running standard deviation.
• Replicated the data pipeline with tokenization, concatenation, 8192-token slicing, seeded global shuffling, and n-gram repetition filtering, all built on Google's Grain library.
• Created checkpoint-conversion and verification tools with logit-parity checks to ensure framework translations preserve model state.

Matching on What Matters: Four Evaluation Surfaces#

Rather than relying solely on training loss curves, the team verified convergence across four independent metrics at six step landmarks spanning 915,000 steps: