Reproducing OLMo 3 7B Pre-training in MaxText: A Case Study of Large-Scale Training on TPUs
以 MaxText 重現 OLMo 3 7B 預訓練：Google Cloud TPU 大規模訓練實戰指南
Google engineers reproduced AI2's OLMo 3 7B from scratch on Cloud TPUs using MaxText, covering Stage-1 pre-training (~5.93T tokens) and Stage-2 mid-training (~100B tokens). The team achieved a near-perfect match with the original PyTorch/GPU reference across both loss curves and downstream metrics (e.g., MMLU, GSM8K). Throughout the replication, they resolved critical bugs in data-sharding (Grain) and optimizer precision, while demonstrating a hardware-software co-design win that boosted throughput by 12.4% by optimizing Head-dim.
Key points
Complete Metric Alignment
Beyond matching loss curves, the downstream task accuracy gap never exceeded ±0.005 macro, proving the high fidelity of the MaxText/TPU stack.
Exposing Data Bugs Hidden as Gains
A double-sharding bug in Grain caused the model to overfit on repeated data, lowering training loss artificially; held-out evaluation successfully exposed it.
12.4% Speedup via Head-Dim Tuning
Changing heads to 16 and dim to 256 perfectly aligned with TPU's 256x256 MXU, yielding a 12.4% throughput boost without changing parameter counts or FLOPs.
Flexible Architecture & Portability
JAX/XLA decouples the recipe from hardware topology. Stage 1 ran on Ironwood, while Stage 2 seamlessly ported to TPU v5p with bit-perfect resumes.
How it works
| 原生配置 (Stock Variant) | MXU 優化配置 (Optimized Variant) | |
|---|---|---|
| Query Heads | 32 | 16 |
| Head Dimension | 128 | 256 |
| MXU Alignment | 未對齊 (50% 陣列閒置) | 完美對齊 (防止閒置週期) |
| Hardware Throughput | 508 TFLOP/s/device (44.2% MFU) | 571 TFLOP/s/device (49.6% MFU) |
| Params & FLOPs | 完全一致 (7.298B Params / 1565 TFLOPs) | 完全一致 (7.298B Params / 1565 TFLOPs) |
Why it matters
Most reproduction efforts stop at comparing training loss curves, which can mask critical bugs. This case study establishes a rigorous standard for cross-framework (PyTorch to JAX) and cross-hardware (GPU to TPU) migration, proving that only held-out metrics can confirm true fidelity. Key technical insights—such as hardware-aligned attention dimensions and the necessity of keeping Adam states in float32—serve as an essential playbook for developers executing stable, cost-effective large-scale LLM pre-training.
Who it affects
- AI Developer
- AI Researcher
- Enterprise Leader
How to use it
- 1Pre-training or fine-tuning LLaMA/OLMo-style LLMs from scratch on Google Cloud TPUs using the open-source MaxText framework.
- 2Configuring exact, stateful resumes in preemptible TPU clusters to guarantee bit-perfect replay of data iterators post-interruption using Grain.
- 3Optimizing Transformer dimension configurations (e.g., attention heads and dimensions) to match target hardware matrix units for hardware efficiency gains.
Limitations & caveats
- Slight discrepancies remained in the first 8k steps of Stage-2 due to random shuffling and data distribution differences between the two independent runs.
- Keeping Adam optimizer moments in bfloat16 silently demotes updates, compounding a +0.93 loss inflation that must be avoided by keeping them in float32.
Related
Minimally Invasive Steering of LMs: Optimizing Rewards Without Quality Degradation
微創型語言模型導向技術：利用 MISVO 在不損害生成品質下優化輸出
This paper introduces MISVO, a minimally invasive steering method that uses local KL geometry to optimize LLM outputs for test-time rewards without parameter updates or quality degradation.
Accelerating MoE Training for Biological Foundation Models with NVIDIA Transformer Engine
NVIDIA Transformer Engine 加速生物基礎模型 MoE 訓練：吞吐量提升達 2.21 倍
This guide demonstrates how to use NVIDIA BioNeMo and Transformer Engine's optimized primitives to overcome MoE training bottlenecks in biological models, boosting throughput by up to 2.21x.
LFM2.5-VL-DSpark: Accelerating Vision-Language Models with Minimal Overhead
LFM2.5-VL-DSpark：以超低開銷將多模態模型推論速度提升達 3 倍
Liquid AI has released a 280M parameter DSpark draft model for LFM2.5-VL-3B, boosting decoding speeds up to 3.13 on-device with only 8.9% parameter overhead.