Google reproduces OLMo 3 7B training on TPUs with MaxText and exposes bugs that only appear in long runs

The Google Developers Blog published a detailed case study this Thursday, September 24, on the full reproduction of OLMo 3 7B's pre-training, an open model from the Allen Institute for AI (AI2), using MaxText, a JAX/XLA training framework for TPUs. This isn't a product announcement: it's a technical account of weeks of real training, complete with the numbers, the charts, and, above all, the two bugs the team found along the way. For anyone building ML infrastructure, this is rare material: most 'model reproduction' posts only show the loss curve hitting the mark. This one also shows where it nearly lied.

The choice of OLMo 3 wasn't accidental. It's a modern 7B model, trained at real production scale (roughly 5.93 trillion tokens over 1.41 million steps), with publicly available data, code, configurations, checkpoints, and training logs, including a reference run published on Weights & Biases. That gave the Google Cloud team something rare in open research: an independent PyTorch/GPU reference against which to test whether MaxText, running on TPU, reproduces not just the architecture but the entire training behavior.

From PyTorch to JAX, with proof of parity

OLMo 3's architecture has some out-of-the-ordinary choices: a 'reordered norm' block, QK-norm, and a 3:1 ratio between sliding-window attention and global attention. Migrating this from PyTorch to JAX and porting it to MaxText wasn't just a matter of running it and comparing the final loss; the team built a conversion checkpoint with logit parity verification. The converted step-0 checkpoint matched the HuggingFace reference at KL ≈ 1.5e-3, which the authors call the 'noise floor between the same model on different frameworks,' and, at a full 8,192-token context in bfloat16, the two agreed on the top-1 token 98.75% of the time.

This kind of verification is what separates a reliable framework migration from one that merely looks like it works. For anyone who has tried porting a model between PyTorch and JAX, or even between versions of the same library, the lesson is direct: comparing only the final output of an isolated forward pass guarantees nothing about production behavior. You need a parity test that catches numerical divergence early, before spending weeks of TPU time training on top of a conversion bug.

The bug that looked like a win

The most interesting part of the post isn't the architecture, it's the demonstration that training loss alone doesn't prove convergence. Starting at around 900,000 steps, MaxText's training loss began falling systematically below the curve published by AI2, and never crossed back. Near 1.25 million steps, the gap reached -0.25 in windows of a few hundred steps. Looking only at that chart, the obvious conclusion would be: MaxText beat the reference.

It didn't. Held-out loss (on C4 data not seen during training) was tied at checkpoints around that stretch (a delta of -0.004 at 1 million steps, +0.003 at the end of stage 1), and accuracy across eight downstream tasks (MMLU, HellaSwag, ARC, OpenBookQA, PIQA, BoolQ, WinoGrande) slightly favored AI2's run at the same point. Training loss dropping without generalization moving is the classic signature of memorization: the model was seeing the same sequences more than once.

The cause was a double-sharding bug in the Grain data loader. OLMo's loader in MaxText was passing ShardOptions(shard_index, shard_count) to Grain's DataLoader while the index sampler was already doing its own sharding internally. The result: with shard_count=32, the data cursor advanced 32 times faster than expected, turning what should have been a clean epoch into resampling with replacement, approximately Poisson(≈1): 37% of the corpus was never seen, 37% was seen once, and 26% was seen twice or more. The total token budget stayed correct (which is why the global loss still tracked AI2's), but the localized repetitions inflated the training metric exactly where they occurred.

The fix was one line: switching to grain.sharding.NoSharding() and letting the sampler be the sole party responsible for sharding. Finding that line required an A/B harness, a unit test that reproduces the divergence whenever shard_count>1, and a new hardware run to validate it. The team let the in-progress run finish as it was (85% complete; the fix doesn't undo data already read, and relaunching would have thrown away 1.2 million steps of computation), documenting that the bug cost zero observable accuracy, while fixing it for future runs.

Checkpoint, resume, and resize without touching the recipe

The second bug was subtler: an off-by-one in resume-step detection. Checkpoint directory N was written after iteration N finished, so the model was restored at step N+1 while the data loader resumed at batch N, retraining one batch and staying permanently one step behind. With both fixes applied (sharding and off-by-one), a controlled checkpoint-and-resume test replicated the uninterrupted run exactly: a delta of 0.000 in logged loss across all 99 steps tested. When a host failure killed the stage-2 run midway through, the resumed run redid 127 steps with a delta of 0.000 in logged loss and perplexity.

This level of exactness matters because a run of 1.4 million steps, spanning weeks, is going to get interrupted. The team uses a resume_until_done loop that automatically resubmits the job after preemption, with checkpoints every 2,000 steps (capping the cost of an interruption to a few minutes of recomputation on the large slice) and a configurable 300-second backoff to avoid exhausting resubmission attempts against Kueue.

The most reusable property of the whole JAX/XLA stack here, according to the post, is that the training recipe is decoupled from the hardware topology: the global batch (512 instances, 4.19 million tokens per step) is fixed, but the number of chips it's distributed across isn't. When the team lost three-quarters of its allocated capacity around step 1.05 million, the run continued on a slice four times smaller, without changing the script (run_olmo3_7b_stage1.sh adjusts the per-device batch to keep the global batch constant), keeping per-device throughput within 1% and maintaining near-100% strong scaling in both directions (from 128 to 512 devices and back).

Switching TPU generations mid-recipe

Stage 2 (mid-training/anneal) was trained on a different TPU generation than stage 1: instead of Ironwood, the same launcher pointed to TPU v5p, changing only the device type, and sustained 57.4% MFU (model FLOPs utilization). In stage 1, on Ironwood, the team reached 44.5% MFU (510-513 TFLOP/s per device) on the original architecture, via offloading collectives (all-gather, reduce-scatter) to SparseCore, v7x-specific XLA flags, extended rematerialization of the attention and MLP projections, and splash attention with 2,048-token blocks.

A detail worth noting for anyone optimizing sharding: in MaxText, sharding topology (pure FSDP vs. combinations of FSDP and tensor parallelism) was practically irrelevant at this scale, with only about 1.5 TFLOP/s of variation between configurations at 128 devices. Intra-chip tensor parallelism (TP=2), on the other hand, was a net loss of 1.6% MFU at reduced batch size and ran out of memory at full batch.

A separate ablation, not used in the official reproduction, tested reshaping attention from 32 heads with dimension 128 to 16 heads with dimension 256, keeping parameters and FLOPs identical. The gain was 12.4% in speed, because head dimension 256 fully utilizes Ironwood's 256×256 MXU, and the loss curve matched the original up to 120 billion tokens (30,000 steps). It's a reminder that, on TPU, the shape of the attention tensor can matter as much as the algorithm.

What's in it for anyone training outside Google

Few teams in Brazil are going to train a 7B model from scratch over millions of steps, but the value of this post isn't in the scale, it's in the verification discipline. Three practices are replicable in any fine-tuning or continued-training project, even on rented GPUs: first, never validate convergence by training loss alone, always cross-check against a held-out set and, if possible, a battery of downstream tasks, because memorization and data-loader bugs hide exactly where training loss seems to improve; second, test numerical parity whenever there's a framework or checkpoint conversion, with an objective metric like KL divergence between output distributions, not just visual comparison of outputs; third, treat resume-from-checkpoint as something that needs automated regression testing, because a slightly wrong resume (like the off-by-one reported here) doesn't break visibly, it silently degrades training.

The code and configurations used in the reproduction, including the olmo3-7b-pt.yml file and the olmo_grain data pipeline, were submitted as public pull requests to the MaxText repository, and the post lists the corresponding issue numbers. That means anyone wanting to run a smaller version of the same experiment, say, a model in the 1B range on a rented TPU v5e via Google Cloud, has a documented starting point, including the two bugs already fixed upstream. Stage 3 (long context, with YaRN) and post-training via SFT and GRPO using Tunix haven't been run by the team yet, so that part of OLMo 3's recipe remains open.