---
title: "Reproducing OLMo 3 7B Pre-training in MaxText: case study of large scale training on TPUs"
description: "The post describes the reproduction of AI2’s OLMo 3 7B pre-training using MaxText on Google Cloud TPUs, covering both stages, matching the reference loss curve and held-out metrics. PyTorch weights were converted to JAX with logit parity, KL divergence 1.5e-3 and 98.75% top-1 agreement, and held-out evaluations caught a data loader bug that made incorrect gains appear. The run demonstrated reliability via exact checkpoint-resume with zero loss delta, resilience to a 75% capacity loss via resize with 1% throughput preservation, and seamless TPU generation change. Performance optimizations achieved 44.5% MFU on Ironwood, and head reshaping from 32 heads by 128 to 16 by 256 yielded 12.4% faster throughput at identical parameters and FLOPs. The post emphasizes that faithful reproduction requires verification on generalization metrics and provides instructions for readers to reproduce the run with open sourced code."
---

# Reproducing OLMo 3 7B Pre-training in MaxText: case study of large scale training on TPUs

[Google](https://yomu.fyi/company/google) · Gagik Amirkhanyan, Ran Ran, Aireen Mei, Matt Davidow, TPU Inference Software Engineering Team, Google Cloud, AI2 Team · Sep 24, 2026

**Type:** Problem & solution

## Summary

The post describes the reproduction of AI2’s OLMo 3 7B pre-training using MaxText on Google Cloud TPUs, covering both stages, matching the reference loss curve and held-out metrics. PyTorch weights were converted to JAX with logit parity, KL divergence 1.5e-3 and 98.75% top-1 agreement, and held-out evaluations caught a data loader bug that made incorrect gains appear. The run demonstrated reliability via exact checkpoint-resume with zero loss delta, resilience to a 75% capacity loss via resize with 1% throughput preservation, and seamless TPU generation change. Performance optimizations achieved 44.5% MFU on Ironwood, and head reshaping from 32 heads by 128 to 16 by 256 yielded 12.4% faster throughput at identical parameters and FLOPs. The post emphasizes that faithful reproduction requires verification on generalization metrics and provides instructions for readers to reproduce the run with open sourced code.

## Context

AI2’s OLMo 3 7B is a fully open language model with open weights, data, and training recipe. The team used MaxText, a JAX/XLA LLM training framework, to reproduce it on Google Cloud TPUs as a validation of the stack against a PyTorch/GPU reference.

## Approach / What changed

The team converted AI2’s step-0 PyTorch weights to JAX Orbax format, implemented the OLMo 3 architecture (reordered norm, QK-norm, 3:1 sliding/global attention) in MaxText, built the olmo\_grain data pipeline matching AI2’s tokenization and repetition filter, and ran stage 1 on Ironwood TPUs with the same recipe but a single cosine LR schedule and the public data mix. They verified with logit-parity checks and held-out evals. For stage 2, they used the same launcher on v5p with a warm-Adam init and Dolmino index. The system handled a resize from 512 to 128 devices and preemption with exact resume.

## Takeaways

- Verification on held-out evaluation metrics caught a data-loader bug that made incorrect gains appear, highlighting the importance of generalization checks over training loss alone.
- Checkpoint-resume in MaxText was proven exact, with a controlled A/B test showing zero loss delta at every step after resume, and surviving a host failure during stage 2.
- A head-reshaping from 32 heads x 128 head dim to 16 x 256, at identical parameters and FLOPs, yielded 12.4% faster throughput because the larger head dim fully utilizes Ironwood’s 256 x 256 MXU.

## FAQ

### How was the PyTorch-to-JAX weight conversion verified?

The converted step-0 checkpoint matched the HuggingFace reference at KL ≈ 1.5e-3, the 'same model, different framework' noise floor, and at full 8192-token context in bfloat16 the two agreed on the top-1 token 98.75% of the time.

### How did the run maintain reliability during failures?

Checkpoint-resume replays the run exactly; a controlled A/B shows Δ = 0.000 at every step after 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.

### What performance gains came from changing the attention head shape?

Reshaping heads from 32 heads x head-dim 128 to 16 x 256, at identical parameters and FLOPs, ran 12.4% faster because head-dim 256 fully utilizes Ironwood's 256x256 MXU, and its loss curve matched the original through 120B tokens (30k steps).

**Tags:** [JAX](https://yomu.fyi/topic/jax), [LLMs](https://yomu.fyi/topic/llm), [MaxText](https://yomu.fyi/topic/maxtext), [Performance](https://yomu.fyi/topic/performance), [TPU](https://yomu.fyi/topic/tpu)

- Source: [Google](https://developers.googleblog.com/reproducing-olmo-3-7b-pre-training-in-maxtext-case-study-of-large-scale-training-on-tpus)
- Source URL: https://developers.googleblog.com/reproducing-olmo-3-7b-pre-training-in-maxtext-case-study-of-large-scale-training-on-tpus
- Ingested by Yomu: 2026-09-24T18:00:12.578Z

[Read original post](https://developers.googleblog.com/reproducing-olmo-3-7b-pre-training-in-maxtext-case-study-of-large-scale-training-on-tpus)