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:
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.
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 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).
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 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.
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:
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):
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.
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:
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.
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.
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:
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.)
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).
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:
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.
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.)
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.
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).
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.
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.
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.
Stage-2 surfaced a resume bug the stage-1 fixes didn’t cover: they made resume exact only at grain_worker_count=1. Stage 2 runs 4 workers (matching AI2), and there the stateless resume (re-derive the data offset from the step) diverges, because a fresh loader started at an offset interleaves the workers’ records differently than the uninterrupted stream. A parity test pins it: stateless resume matches at 0/2 workers, mismatches at 4; Grain’s own iterator-state checkpoint is bit-exact at every worker count.
So for stage-2 we wired olmo_grain through Grain’s GrainCheckpointHandler: the checkpoint now carries an iter item alongside the model items, and a resume restores weights, optimizer state, and the exact data-iterator position together, atomically. On a contended cluster where a 47,684-step run will be preempted many times, this is what keeps the full run a single clean epoch instead of a noisy resample-with-replacement.
The proof arrived unplanned. Sixteen hours in, a host failure killed the jobset at step 19,627 (loss smooth right up to it: an infra blip, not a divergence). The run resumed from the step-19,500 checkpoint and re-trained the 127 lost steps, about six minutes of recompute on a 47,684-step run. Because those steps were already in the logs before the crash, the preemption handed us a free paired diff, and the re-trained loss matched the original exactly: Δ = 0.000 in both loss and perplexity at all 127 steps, against the ~0.01–0.4 scatter a broken resume produces. And it did so at grain_worker_count=4, the exact configuration where stateless resume fails. That an interrupted run replays exactly isn’t luck; it’s three deterministic layers compounding: XLA’s compiled TPU program, the Grain iterator state, and the Orbax restore.
One number looked wrong early on: with no pad masking, our stage-2 training loss spiked to ~7.5 on the olmOCR PDF shards (<|pad|>, token 100277, ~2% of mid-training tokens), while AI2’s curve there is smooth (max ~1.96). The cause is upstream: our stage-1 n-gram repetition filter masked all-pad windows, so our model never learned to predict <|pad|>, and scores ~17 loss when forced to, while AI2’s loss is effectively pad-penalty-free. The faithful choice for our pad-naïve model is to mask the pad token in the loss (olmo_pad_token_id=100277); a paired LR=0 eval on identical batches confirms the models are otherwise equivalent (Δ −0.001). “Match the reference recipe” sometimes means accounting for an upstream-stage difference, not blindly copying the current-stage config.
The full 100B run finished at step 47,683 with the LR decayed to 0, having auto-resumed through every interruption. The question that matters isn’t the training curve; it’s whether the finished model matches AI2’s mid-trained checkpoint on data it never trained on.
We evaluate our final checkpoint against AI2’s Olmo-3-1025-7B stage-2 step-47684 with one eval stack for both: held-out C4 loss, multi-domain perplexity (Paloma-style), and the 8-task lm-eval suite.
The loss-side metrics all agree at a steady ~+0.006 nat: the same tiny offset on training CE, held-out C4, and the perplexity mean. That uniformity points to cross-framework data ordering rather than a bug; a real bug blows out one domain, not all of them equally. Per-domain perplexity confirms it: wikitext103 ties (−0.0005), and the widest domain is wikipedia at +0.0135 (Appendix M).
Crucially, that loss offset does not become a capability gap. Downstream accuracy is a tie: macro −0.0023, with the eight tasks scattering in both directions inside eval noise (per-task metric: lm-eval default, acc_norm where defined, else acc):
And it’s not just the finish line. Evaluating both stacks at three landmarks across the anneal shows the held-out gap stays flat and tiny the whole way: C4 holds ~+0.006 nat at all three landmarks, 6-domain perplexity stays within +0.0027–0.0065, and both models improve in lockstep as the LR decays to 0:
So MaxText reproduced OLMo-3 7B’s mid-training end to end: a ~+0.006-nat loss offset consistent with data-order noise, and downstream accuracy indistinguishable from AI2’s (macro −0.0023, inside eval noise). The anneal also moved capability the way AI2’s did: on a like-for-like metric basis, MMLU 5-shot jumps 0.605 → 0.648 for us and 0.605 → 0.650 for AI2, with the 8-task macro up ~+0.6 points (AI2: +0.9). (A metric note: the stage-1 tables report plain acc, while the stage-2 table above uses lm-eval’s default, acc_norm where defined, so the two tables’ macros are not directly comparable; the like-for-like numbers in this paragraph are computed on matching bases.)
A multi-week, cross-stack reproduction teaches more than the final number. What we’d hand the next team:
Stages 1 and 2 are done. Still ahead:
The recipe, launcher, and conversion tooling live in MaxText. The shape of a run:
# Convert AI2's step-0 weights to Orbax, then:
export OLMO_INDEX_PATH=/path/to/olmo_index_seq8192.json
export LOAD_PARAMETERS_PATH=gs://<your-bucket>/olmo/checkpoints/stage1-step0/0/items
bash src/maxtext/trainers/pre_train/scripts/olmo/xpk_olmo3_7b_stage1.sh submit
# Drive the full run, auto-resubmitting on preemption:
STEPS_OVERRIDE=1414078 bash src/maxtext/trainers/pre_train/scripts/olmo/xpk_olmo3_7b_stage1.sh resume_until_doneStage 2 (mid-training) reuses the same pipeline, pointed at the Dolmino index, from a warm-Adam init:
# 1. Build the warm-Adam stage-2 init from the stage-1 final full-state checkpoint
# (keeps params + Adam mu/nu, resets the loop step + LR-schedule counter):
python3 -m maxtext.utils.olmo3_build_stage2_init # src/dst paths in the script header
# 2. (Once) build the Dolmino 100B index from the public mix file:
python3 tools/data_generation/build_olmo_npy_index.py \
--mix-file OLMo-midtraining-mix-0625-100B.txt --gcs-base gs://<bucket>/<dolmino-prefix>/ \
--tokenizer allenai/dolma3-tokenizer --sequence-length 8192 \
--output olmo_midtraining_index_seq8192.json
# 3. Launch the full 100B anneal. The launcher already defaults to the AI2-faithful
# recipe (LR 2.0712e-4 → 0 linear, warmup 0, 47,684 steps, batch 256, seed 1337,
# pad-mask 100277) and resumes exactly (weights + optimizer + data iterator):
export OLMO_INDEX_PATH=/mount/olmo_midtraining_index_seq8192.json
export LOAD_FULL_STATE_PATH=gs://<bucket>/olmo/checkpoints/stage2-init/checkpoints/0/items
STEPS_OVERRIDE=47684 bash src/maxtext/trainers/pre_train/scripts/olmo/xpk_olmo3_7b_stage2.sh resume_until_doneThe launcher pins every hyperparameter from the recipe table, exports the Ironwood XLA flag set, and computes per-device batch from the global device count. The comparison tooling (tools/wandb_csv_to_tensorboard.py, tools/compare_loss_curves.py, tools/eval_lm_loss.py) is what produced every table in this post: point it at two TensorBoard dirs and it prints the deltas.
The bottom line: a PyTorch-on-GPU pre-training recipe reproduces faithfully in JAX-on-TPU, but only “faithfully” if you measure generalization, not just training loss. The single most important methodological choice we made was to verify on held-out eval at every landmark. It caught a data bug masquerading as a win, and it’s what lets us say “reproduced” with a straight face.
