Reproducing OLMo 3 7B Pre-training on TPUs with MaxText
A case study in large-scale language model training reproducibility
1OLMo 3's unique openness
The Allen Institute's OLMo 3 stands out as one of the few fully open frontier-class language models, providing not just weights but complete training artifacts including data, code, configurations, checkpoints, and evaluations. This comprehensive transparency enabled Google's team to conduct a rigorous reproduction study using their MaxText framework on TPUs, comparing against AI2's original PyTorch/GPU implementation across multiple validation surfaces.
2Reproduction methodology
The team reproduced OLMo 3's 7B parameter model training in two stages: initial pre-training on ~5.93 trillion tokens and mid-training annealing. They started from AI2's initial weights but proved the match wasn't initialization-dependent by also testing from random initialization. The MaxText implementation used JAX/XLA on TPUs while meticulously mirroring the original recipe's key components - model architecture, optimizer settings, and data pipeline - with only two intentional simplifications to the learning rate schedule and data mix.
3Architecture specifics
OLMo 3's architecture features a 32-layer, 4096-dimension dense transformer with several non-standard choices including reordered norm blocks, QK-norm, and a 3:1 mix of sliding-window/global attention. The 7B parameter model was exactly replicated in MaxText's configuration, maintaining identical computational characteristics (7.298B parameters, 1565 TFLOPs/step) while adapting the implementation to TPU hardware constraints through JAX/XLA optimizations.
4Training curve matching
Despite different frameworks (PyTorch vs JAX) and hardware (GPU vs TPU), MaxText's training loss closely tracked AI2's published curve throughout the ~1.41 million training steps. The match held within ±0.012 through 800k steps, with MaxText eventually running slightly below the reference due to a data loading bug discovered later. Crucially, validation on held-out metrics confirmed the implementations were functionally equivalent beyond just training loss.
5Validation surfaces
The team established reproduction fidelity through four independent evaluation metrics at six step checkpoints spanning 915k steps. These included language modeling loss on C4, downstream task accuracy (MMLU), token-level KL divergence, and sequence-level perplexity. All metrics showed the implementations were interchangeable, with downstream accuracy gaps never exceeding ±0.005 and randomly flipping direction, indicating only run-to-run variation.
6Data pipeline bug discovery
A critical bug emerged where MaxText's data loader was advancing 32x too fast due to double-sharding, causing some sequences to repeat while others were missed. This created a memorization effect where training loss decreased without corresponding generalization improvements. The team diagnosed this through careful validation showing tied C4 loss despite diverging training curves, then implemented fixes with regression tests to prevent recurrence.
7Hardware optimization insights
The reproduction revealed significant hardware-software co-design opportunities. By adjusting the attention head configuration from 32×128 to 16×256 dimensions, the team achieved a 12.4% throughput increase (571 vs 508 TFLOP/s/device) without changing parameters or FLOPs. This better aligned tensor shapes with TPU's 256x256 MXU systolic arrays, eliminating idle compute cycles during attention operations.
8Numerical precision pitfalls
One subtle but impactful discovery involved optimizer numerical precision. Using bfloat16 for weight updates silently degraded performance by accumulating rounding errors during the warmup phase, adding +0.93 to the loss over 1000 steps. Maintaining float32 precision for Adam's momentum variables while using bfloat16 elsewhere preserved numerical stability, closing this gap without computational overhead.
9Cross-hardware portability
The JAX/XLA stack demonstrated remarkable hardware flexibility when stage 2 training moved from Ironwood to TPU v5p chips without recipe changes. The identical launch script achieved 57.4% MFU on v5p (263 TFLOP/s/chip) versus 44.5% on Ironwood, showing how decoupling the recipe from topology enables efficient utilization of whatever hardware is available. This portability is crucial for long-running training jobs.
10Training cost breakdown
Stage 1 pre-training consumed ~77k Ironwood chip-hours (~3,200 chip-days), with checkpointing and evaluation overhead. Performance optimizations saved ~36k chip-hours compared to initial estimates. Stage 2 annealing was more efficient at ~5k v5p chip-hours (~39 hours on 128-chip slice). The total compute investment validated both the reproducibility of OLMo's training and MaxText's efficiency on TPU hardware.
11Reproduction significance
Successfully matching OLMo 3's training on different hardware/software validates MaxText as a faithful LLM training framework. More importantly, it demonstrates that modern large-scale training recipes can be precisely reproduced across frameworks when all components are properly specified and validated. This work sets a benchmark for open, reproducible LLM development and provides concrete methodology for cross-framework validation.
12Future implications
The case study highlights critical considerations for large-scale training: the importance of validation beyond loss curves, hardware-aware architecture tuning, numerical precision management, and fault-tolerant data loading. These lessons directly inform best practices for both framework developers and ML practitioners undertaking similar reproduction efforts or original training runs at scale.
Original article: Google Developers Blog — Reproducing OLMo 3 7B Pre-training on TPUs with MaxText