MaxTextでOLMo 3 7Bの事前学習を再現：TPUを用いた大規模学習のケーススタディReproducing Olmo 3 7B Pre-training in MaxText: case study of large scale training on TPUs

Reproducing OLMo 3 7B

The MaxText team reproduced AI2’s OLMo 3 7B language model from scratch on Google Cloud TPUs using JAX/XLA. The case study matters because it tested whether large-scale training could match the original PyTorch-on-GPU reference while remaining portable and dependable across infrastructure changes.

The reproduction covered both pre-training and mid-training. Rather than relying only on training loss, the team compared the implementations through held-out evaluations.

Matching the GPU reference

Across all held-out evaluations, the MaxText implementation precisely matched the original PyTorch-on-GPU reference during the pre-training and mid-training stages. This established that the JAX/XLA and TPU implementation reproduced the reference results rather than merely reporting a lower training loss.

The implementation also achieved up to 57.4% Model Flops Utilization. That result accompanied the held-out evaluation match, providing both a training-efficiency measurement and an evaluation-based comparison with the GPU reference.

Infrastructure portability

The training run demonstrated robust infrastructure portability. It survived cluster resizes in the middle of the run and shifts between TPU generations without requiring alterations to the training recipe.

These events tested the implementation under changing infrastructure conditions while preserving the same recipe. The successful continuation across mid-run resizing and cross-generation TPU shifts was part of the case study’s evidence for portable large-scale training.

Comprehensive held-out validation

Held-out validation also exposed a silent data-loader memorization bug. The bug artificially depressed training loss, creating what would otherwise have appeared to be a performance win.

Because the issue was silent and made the loss look better, training loss alone would not have revealed that the result was misleading. Comprehensive held-out evaluations caught the memorization problem and prevented the lower loss from being treated as a genuine improvement. The exercise therefore demonstrated that reproducing large-scale training requires validation beyond the training metric itself.