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

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.

Takeaways

1. 1

   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.

2. 2

   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.

3. 3

   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