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

SEPT. 24, 2026

Ran Ran Staff Software Engineer

OLMo 3, developed by the Allen Institute for AI (AI2), is a state-of-the-art, fully open language model trained with a modern architecture and a multi-stage training recipe. To evaluate the capabilities of MaxText on Google Cloud TPUs, our team set out to reproduce AI2’s OLMo 3 7B from scratch. We chose OLMo 3 because it combines three properties that rarely appear together. It is a strong, modern 7B model trained at real production scale. AI2 exposes nearly the complete model flow, including data, code, configurations, checkpoints, logs, and evaluations. And finally, it gives us an independent PyTorch and GPU reference against which we can test MaxText and TPUs.

We reproduced AI2’s OLMo 3 7B in MaxText on Google Cloud TPUs, both the stage-1 pre-training and the stage-2 mid-training anneal, and proved the match on held-out metrics, not just the loss curve:

The main highlights, each covered in detail later in the post:

• PyTorch → JAX model conversion. OLMo 3's architecture (reordered-norm block, QK-norm, 3:1 sliding/global attention) ported to MaxText and verified with a logit-parity check: the converted step-0 checkpoint matches the HuggingFace reference at KL ≈ 1.5e-3, the "same model, different framework" noise floor, and at the full 8192-token context in bfloat16 the two agree on the top-1 token 98.75% of the time.

• Verification that catches real bugs. Held-out evals caught a data-loader bug that made MaxText look like it was beating the reference; the gain was memorization.

• Reliability over a multi-week run. Checkpoint-and-resume replays the run exactly: a controlled A/B shows Δ = 0.000 at every step after a resume, and when a host failure killed the stage-2 run mid-flight, the resumed run re-trained 127 steps at Δ = 0.000 in logged loss and perplexity.

• Resizing the training job mid-flight. At step ~1.05M we lost three quarters of our capacity; the run resumed on a slice one quarter the size with no recipe change (same script run_olmo3_7b_stage1.sh scales per device batch size to keep GBS constant), per-device throughput preserved within 1% (≈100% strong scaling, measured in both directions).

• Changing TPU generation mid-recipe. Stage 2 pointed the identical launcher at v5p instead of Ironwood, changing only the device type, and sustained 57.4% MFU.

• Performance work that paid for itself. 44.5% MFU on Ironwood at 7B via SparseCore collective offload, remat tuning, and optimal sharding, roughly a third of the compute budget bought back.

• Co-design for TPU: faster at the same quality. Reshaping attention from 32 heads × head-dim 128 to 16 × 256, at identical parameters and FLOPs, runs +12.4% faster because head-dim 256 fully utilizes Ironwood's 256×256 MXU, and its loss curve matches the original through 120B tokens (30k steps). This was a side ablation; the reproduction kept the original architecture.

Starting from AI2’s step-0 PyTorch weights and the same core recipe, the MaxText run tracks AI2’s published loss curve over the full ~5.93T-token / 1.41M-step budget and lands on top of it at the end of stage-1. We even simplified two recipe details (a single cosine LR schedule where AI2 stitched two, and the publicly released data mix; see the recipe below), and the match held anyway. The rest of this post is how each of these was built, measured, and, in one instructive case, nearly faked.

Why reproduce OLMo 3?

OLMo 3 is one of the few genuinely open frontier-class language models: open weights, open data, and a fully specified training recipe with a public reference run on Weights & Biases. Matching that independently trained run, on held-out metrics rather than just the loss curve, is strong evidence that the MaxText stack (optimizer, loss, data pipeline, numerics) is faithful, not just “looks like it’s training.”

MaxText is a JAX/XLA LLM training framework built for TPUs. The question we set out to answer: can a PyTorch-on-GPU recipe be reproduced faithfully in JAX-on-TPU, matched on the metrics that matter rather than bit-for-bit, and how do you prove it?

OLMo 3’s recipe is a 3-stage curriculum: general pre-training, mid-training (annealing), and long-context adaptation. This post covers stage 1 (the ~5.9T-token pre-training run) and stage 2 (mid-training), both trained end to end and matched against AI2’s references. Stage 3 and post-training (SFT/RL via Tunix) are recipes we’ve written but not yet run.

OLMo-3 pre-training curriculum: stage 1 (Ironwood) and stage-2 anneal (v5p) reproduced in this post; stage 3 long-context (seq 65k, YaRN) and post-training (SFT then GRPO via Tunix) next The OLMo-3 curriculum. Stages 1 and 2 are reproduced in this post; stage 3 and post-training are next.

The recipe

OLMo 3 7B is a 32-layer, 4096-dim dense transformer with a few non-standard choices: a “reordered norm” block, QK-norm, and a 3:1 mix of sliding-window and global attention. The MaxText config (olmo3-7b-pt.yml, used for stage 1 and 2) matches it exactly:

The training recipe mirrors OLMo-core’s pretrain-1.py; the knobs that have to match for the curves to line up:

We started training from AI2’s step-0 PyTorch checkpoint, converted to Orbax, so MaxText begins from the exact same weights as the reference. The conversion itself was the first checkpoint: a forward pass on the converted weights matched the HuggingFace reference at KL ≈ 1.5e-3 with 9/10 top-10 token overlap, the “same model, different framework” noise floor.

Does the match depend on inheriting AI2’s initialization? Apparently not. As an independent check we also trained a run from MaxText’s own random init for ~50k steps (3.5% of the horizon); its training loss tracked AI2’s published curve closely, running a touch below it. That’s a training-loss spot check, not a full replicate, but it suggests the match doesn’t hinge on starting from AI2’s weights.

The data pipeline mirrors OLMo-core exactly: tokenize and concatenate all documents (EOS between), slice into non-overlapping 8192-token instances, globally shuffle the index with a fixed seed, and apply an n-gram repetition filter that masks instances with >32 repeated n-grams. MaxText’s dataset_type=olmo_grain (built on Grain) implements this.

Two deliberate divergences from AI2’s run. (1) LR schedule: AI2 originally planned ~5T tokens and extended the run mid-flight to a final horizon of ~5.93T, so its LR trace stitches two cosine curves (visible in its public WandB run); we ran a single cosine over the full horizon. (2) Data: we train on the publicly released OLMo-3 mix, which omits a small fraction (<0.5% of the token budget, mostly s2pdf shards absent from the released file list) that AI2’s internal run saw. Both are simplifications we chose, not accidents, and MaxText still matches on every held-out surface. This is also why we say “reproduced to within run-to-run noise,” not bit-for-bit (see the KL analysis in §G).

What we had to build

OLMo 3 wasn't in MaxText when we started; the reproduction added, and upstreamed, everything below. "Reproduce it yourself" at the end of this post is config, not code.

• The model itself: the reordered-norm block, QK-norm, and the 3:1 sliding/global attention pattern (#3004, #3112).

• A skip-step optimizer matching OLMo-core's semantics down to the Bessel-corrected running std (skip at 6σ over a 128-step window) (#3490).

• z-loss (#3211) and per-parameter weight-decay masking so embeddings can be excluded (#3280).

• The olmo_grain data pipeline (#3749): random-access reads of pre-tokenized shards, a seeded global index shuffle with a fingerprint guard against silent data swaps on restart, the n-gram repetition filter, and (from stage 2) Grain iterator-state checkpointing.

• HF↔Orbax checkpoint conversion with a logit-parity check, the tool behind every "framework noise floor" number in this post (#3112, #3832).

• The stage-1 and 2 launcher (env-driven run script + XPK wrapper with submit / monitor / resume_until_done) (#3886).

• TensorBoard parity tags (optim/step_skipped, perf/total_tokens) so every metric on AI2's W&B dashboard has a MaxText counterpart to compare against.

Does it converge?

The headline is a single overlay: MaxText’s stage-1 lm_loss vs AI2’s published WandB curve, step-aligned and binned to 2k-step means. Through ~800k steps the two track within ±0.012; from ~0.9M MaxText edges below AI2 and never crosses back, the first sign of the data bug dissected in the next section.

MaxText vs AI2 stage-1 loss over 1.41M steps, with the 18k-step-smoothed gap in a lower panel Top: the curves are indistinguishable at this scale until the tail. Bottom: the gap stays inside ±0.012 through ~800k, tilts negative from ~0.9M as data-bug repeats accumulate, and dives past 1.25M, reaching −0.22 in raw 2k-step bins before the 18k-step smoothing both panels are drawn with. None of this reaches held-out loss or downstream accuracy, as the next sections show. (Per-landmark table: Appendix A; full curves committed as olmo_stage1_loss_curve.tsv.)

But a loss curve alone is a weak proof: two runs can match on training loss and diverge on everything you’d actually care about. So we verified convergence on four independent surfaces at six step landmarks spanning 915k steps:

• Held-out C4 lm_loss: forward-only eval on 16M tokens of held-out C4-en, identical batches both runs.

• 8-task lm-eval-harness suite: MMLU, HellaSwag, ARC-easy/challenge, OpenBookQA, PIQA, BoolQ, WinoGrande.

• Multi-domain held-out perplexity: a Paloma-style sweep across web, news, encyclopedic, and mixed domains.

• Token-level KL: next-token distribution distance on identical inputs.

How we measured. All evals run on step-aligned checkpoint pairs: the live MaxText run vs the AI2 checkpoint at the same step, i.e. its public HuggingFace revision allenai/Olmo-3-1025-7B@stage1-step{N} converted to Orbax (the conversion reproduces the PyTorch reference to KL ≤ 1.8e-3, the framework noise floor). lm-eval uses the standard lm-eval-harness (5-shot MMLU, defaults elsewhere); σ is the per-task harness stderr, combined in quadrature for deltas.

At end of stage-1, every surface agrees the two recipes are interchangeable:

The fourth surface, token-level KL, is the one number that isn’t tiny: mean 0.389 nats on identical inputs, ~200× the framework noise floor. That’s expected for two independent runs of the same recipe (same aggregate skill, different allocation of probability mass), and it’s why we say “run-to-run noise,” not bit-for-bit; the breakdown is in §G.

And the downstream-accuracy gap never exceeds ±0.005 macro across all six landmarks, with the sign flipping four times, exactly the random walk you’d expect from two faithful runs differing only in RNG and numerics (per-landmark table in Appendix D):

8-task macro accuracy, MaxText vs AI2 across six landmarks, with per-landmark delta Downstream capability is interchangeable at every landmark. Unlike training loss, accuracy never diverges monotonically; the delta random-walks inside ±0.005 and ends at +0.0002.

The bug that looked like a win

Here’s where it gets interesting. From ~0.9M steps MaxText’s training loss edged below AI2’s and never crossed back; past ~1.25M it pulled clearly under, by −0.06 on average and by as much as −0.25 in a few hundred-step stretches. Watching only the training-loss overlay, you’d conclude MaxText had pulled ahead.

It hadn’t. Held-out C4 loss at the bracketing checkpoints was tied (Δ −0.004 at 1,000k, +0.003 at end of stage-1), and downstream accuracy at 1,350k slightly favored AI2. Training loss was dropping while generalization didn’t move. That’s the signature of memorization: the model was seeing some sequences more than once and scoring low loss on the repeats.

Top: MaxText training loss dives to 1.63 while AI2 stays flat. Bottom: held-out C4 delta stays near zero throughout The whole story in one figure. Top: in the 1.24M–1.41M window MaxText’s training loss repeatedly plunges to 1.63 on 2k-step means (Δ −0.22; −0.25 in hundred-step stretches) while AI2’s stays flat. It looks like a runaway win. Bottom: the training-loss Δ (blue) drifts negative, but held-out C4 loss (green diamonds) never leaves the ±0.02 band. Training loss dropped; generalization didn’t.

The cause was a double-sharding bug in the Grain data loader. MaxText’s OLMo loader passed ShardOptions(shard_index, shard_count) to the Grain DataLoader while the index sampler was already sharding internally. Grain’s shard_options doesn’t just record metadata; it re-strides the sampler’s index stream. With shard_count=32, the data cursor advanced 32× too fast, so stage-1 stopped being one clean epoch and became a Poisson(≈1) resample-with-replacement: roughly 37% of the corpus never seen, 37% seen once, 26% seen twice or more. The token budget was unchanged (~5.9T real tokens), which is why the loss still tracked AI2 globally, but the repeated instances deflated training loss exactly where they recurred.

Two lessons came out of this:

• Training loss is not a convergence proof. The only reason we didn’t ship a false “MaxText beats the reference” claim is that we’d committed to held-out eval at every landmark. The memorization dip is invisible on held-out C4 and on all 8 downstream tasks.

• Honest reproductions need a bug budget. The fix (grain.sharding.NoSharding(), letting the sampler own all sharding) is one line. Finding it took an A/B harness, a unit test that reproduces the divergence at shard_count>1, and a hardware re-run to validate.

While validating the fix we found a second, independent bug: an off-by-one in resume-step detection. The checkpoint directory number is N, but the train loop writes dir N after iteration N completes, so the model restored to step N+1 while the data loader resumed at batch N, re-training one batch and then running permanently one step behind. With both fixes, a checkpoint-and-resume run replays the uninterrupted run exactly: Δ = 0.000 in logged loss at all 99 steps. (That A/B ran single-worker data loading; the multi-worker case surfaced in stage 2, where we closed it; see Stage 2.) Both bugs have regression tests that fail on the old code and pass on the fix.

We let the in-flight stage-1 run finish as-is: it was 85% done, the fix can’t un-scramble already-read data, and a relaunch would forfeit ~1.2M steps of compute. The verification above shows the bug cost zero observable accuracy; the fix is for future runs.

Performance and scale on Ironwood

Reproducing the math is half the job; the other half is making it fast, and keeping it fast when the cluster shifts under you. Over the weeks the 1.4M-step run took, the job was preempted, rescheduled, and resized more than once, and the stack had to absorb all of it without touching the recipe.

Squeezing out MFU

On Ironwood at 7B, per-device batch 4, we landed at 44.5% MFU (510–513 TFLOP/s/device) for the stock architecture (“variant D,” our label from the ablation sweep in Appendix J). A shape-only head-dim change clears 49%; see Head-dim below. What moved the needle, in order of impact:

• Ironwood XLA flags + SparseCore offload: offloading collectives (all-gather, 2D all-gather, reduce-scatter) to the SparseCore, plus a set of v7x-specific XLA flags, took us from 41% to 44.5% MFU, loss-neutral. (Full flag list in Appendix H.)

• Extended rematerialization: checkpointing the attention and MLP projections (qkv_proj, q/k/v_proj, out_proj, mlpwi_0, mlpwo, context) fit the activation-memory budget; adding one more (mlpwi_1) overflowed HBM by 18 GB.

• Splash attention + Tokamax with 2048-token blocks.

• Sharding axis is irrelevant at this scale: pure FSDP, 4-FSDP×32-DP, and 8-FSDP×16-DP were all within ~1.5 TFLOP/s on 128 devices; pure FSDP wins on simplicity. Intra-chip tensor parallelism (TP=2) was a net loss: −1.6% MFU at half batch, OOM at full batch (FSDP=64 doubles per-chip weight state).

Scaling up and down, and why it was free

The single most useful property of the JAX/XLA stack here is that the recipe is decoupled from the topology. The global batch (512 instances, 4.19M tokens/step) is fixed; the number of chips it's spread over is not. (A unit note: Ironwood packs two JAX devices per chip, so the 64-chip 4×4×4 slice exposes 128 devices; we quote both.)

• Scale-up validation: going from a 128-device slice to a 512-device slice (4×) at the same global batch gave a 3.99× aggregate-throughput increase, ≈100% strong scaling, in a 1000-step test. Strong scaling is the hard direction: each device now does a quarter of the work per step, while the collectives span 4× as many devices, so there is less compute available to overlap more communication. SparseCore offload kept that communication off the critical path anyway.

• Scale-down in production: at step ~1.05M we lost three quarters of our capacity, and the run resumed on a 128-device slice (one quarter the size) at the same global batch, with no recipe change. Per-device throughput was preserved (~510–513 TFLOP/s/device on both); only wall-clock per step changed (0.76 s → 3.05 s, the expected 4×).

Left: aggregate throughput scales 3.99x from 128 to 512 devices. Right: per-device TFLOP/s preserved across the resize

This is what lets a long run survive a contended cluster: take whatever capacity is free, keep the math identical. Stage 2 pushed the same idea across TPU generations (see below).

Auto-resume: surviving a multi-week run

A 1.4M-step run will be interrupted. We drive it with a resume_until_done loop that auto-resubmits on preemption and resumes from the latest Orbax checkpoint:

• Checkpoint every 2000 steps, so a preemption costs at most ~2000 steps of recompute (minutes on the large slice). And give the resubmit loop a real backoff: an early version exhausted MAX_RETRIES=50 against Kueue back-pressure; a configurable RETRY_BACKOFF_SECONDS (default 300 s) let it ride out multi-hour scheduling gaps.

• TensorBoard on GCS is the source of truth: kubectl logs only sees the current pod’s history. Every table in this post was generated from GCS-persisted TB events, not live logs.

• Resume has to be exact, or it silently corrupts the run. A resume that reads the wrong data or restarts one step off looks fine on the loss curve but isn’t the run you think it is; that’s exactly the two data-loader bugs above. After the fixes it replays exactly, and the paired A/B below is the proof.

Buggy resume scatters around the continuous run; fixed resume is exactly zero at every step A controlled A/B on 128 devices (64 chips): resume from a checkpoint at step 100, compared to the uninterrupted run. The off-by-one bug (red) desyncs data from parameters (mean |Δ| 0.044, max 0.22) and never reconverges. The fix (green) is exactly 0.000 at all 99 steps. A resume bug is invisible on a normal loss curve; you only see it in a paired diff like this.

Hardware-Software Co-design: A Free 12% Speedup

Alongside the reproduction, we ran an architecture ablation that turned into a major win for hardware-software co-design. OLMo-3 7B ships with a stock configuration of 32 query heads × 128 head-dim. Because num_heads × head_dim = emb_dim = 4096, we can trade heads for width by changing this to 16 heads × 256 head-dim. This architectural adjustment maintains an identical 7.298B parameters and 1565 TFLOP/step, but alters the tensor shape to align beautifully with the underlying hardware.

Head-dim reshape: 44.2% to 49.6% MFU (508 to 571 TFLOP/s per device) at identical params and FLOPs

Because Ironwood's Matrix Multiply Unit (MXU) is a 256x256 systolic array, the standard head-dim of 128 leaves half of the array idle during the attention QK matmul. Reshaping to a head-dim of 256 perfectly aligns the tensor dimension to 256 with the hardware, entirely preventing idle compute cycles. This yields a +12.4% throughput increase (571 vs 508 TFLOP/s/device, or 49.6% vs 44.2% MFU) while keeping parameters and FLOPs completely identical. This is a free speedup that is highly worth implementing before committing a long run to the stock configuration. (loss curve and seed caveats are discussed in Appendix L.)

A gotcha worth its own paragraph

Optimizer dtype is silent and expensive. Setting weight_dtype=bfloat16 silently demoted Adam’s m/v moments via mu_dtype inheritance, adding +0.93 to the loss over 1000 steps: bf16’s ~3-digit mantissa drops a fraction of every tiny early-warmup update, and it compounds. Leaving weight_dtype=float32 (the default) collapsed the gap 30×. This was the single biggest “why doesn’t it match” moment of the project.

Compute

Stage-1 cost ~77k Ironwood chip-hours of step compute (~3,200 chip-days), with checkpointing, eval, and restart ramp on top. Chip-hours is the unit that doesn’t move: chip-seconds per step are slice-independent (0.76 s × 256 chips ≈ 3.05 s × 64 chips ≈ 195 chip-s), while wall-clock depends on the slice. The bulk ran on a 4×8×8 slice (256 chips / 512 devices), where 77k chip-hours is ~12.5 days-equivalent; with the post-preemption stint on the 64-chip slice and time spent queued, calendar time ran to a few weeks. At our pre-run 30%-MFU budget the same tokens would have needed ~50% more chip-time (~113k chip-hours on the same step-time accounting; Appendix I’s 6·N·D planning row reads ~100k), so the perf work bought back roughly a third. Stage 2 was comparatively cheap: ~5k v5p chip-hours (~39 h on a 128-chip v5p-256; details in Stage 2 and Appendix I).

Stage 2: Mid-training (annealing)

With stage-1 matched, we moved to stage 2: mid-training, the final decay of the warmup-stable-decay (WSD) schedule. The stage-1 model is annealed on the Dolmino 100B mix (high-quality math, code, reasoning, and curated web) while the learning rate decays linearly from 2.0712e-4 to 0. This is where OLMo-3’s high-quality data turns into capability gains, so a faithful stack has to match it too. We verified the recipe against AI2’s midtrain reference (run zxv811e1, generated by OLMo-core’s OLMo-3-1025-7B-midtrain.py); every hyperparameter matches:

Warm-Adam init. AI2 re-warms instantly from stage-1’s final 3e-5 to 2.0712e-4 with no warmup and load_optim_state=True; the loaded Adam second moment is what plausibly absorbs that jump. We reproduce it with a one-off checkpoint surgery (keep params + mu/nu, zero the loop step and the LR-schedule counter), so the run restarts the schedule at the peak with warm moments, matching AI2’s init config. Step-0 loss was ~1.53, not a cold-start spike.

Different TPU generation, same launcher: 57.4% MFU on v5p

Stage 2 also moved hardware generations. Ironwood capacity was committed elsewhere, so we pointed the identical launch script, Ironwood-tuned XLA flags and all, at a TPU v5p slice (v5p-256, 128 chips), changing only the XPK device type. With no v5p-specific tuning it landed at 57.4% MFU (263 TFLOP/s/chip median, of v5p’s 459 peak), higher than stage-1’s 44.5% on Ironwood, because a 7B model saturates the older chip more easily than a part with 5× the peak FLOPs. And it stayed there: per-chip throughput sat between 263.0 and 263.9 TFLOP/s from the 25th to the 90th percentile of the entire run, a 0.4% spread over 47,684 steps. Together with the stage-1 resize, that’s the portability story in full: the recipe is decoupled from both slice topology and TPU generation. You run on whatever capacity is free.

Does it converge?

Over the full 47,684 steps the training-loss gap vs AI2 is +0.0044 overall, carried almost entirely by the first ~8k steps. That early gap isn’t a recipe mismatch: early on, the two shuffles have trained on mostly different data. Two random 8k-step prefixes of the 12.2M-instance mix share only ~17% of their instances (just ~2% at 1k steps, where the gap peaks). Once coverage overlaps, the gap washes out: every 4k-step window past step 12k sits between +0.0000 and +0.006, and the back third of the run averages +0.0007, a tie. One measurement asymmetry to note: our CE masks <|pad|> (next section) while AI2’s includes it at near-zero cost, a bias that pushes AI2’s curve down, so +0.0044 is if anything an upper bound on the like-for-like gap.

A stage-2-only resume bug: checkpoint the data iterator too

What’s next

Shell

Copied