# 在 MaxText 中用 Google Cloud TPU 复现 AI2 的 OLMo 3 7B 预训练

- 来源：Google Developers Blog（RSS）
- 作者：Gagik Amirkhanyan
- 发布时间：2026-09-24 23:13
- AIHOT 分数：48
- AIHOT 链接：https://aihot.news/items/cmufoas7707oyro8w46swq4rx
- 原文链接：https://developers.googleblog.com/reproducing-olmo-3-7b-pre-training-in-maxtext-case-study-of-large-scale-training-on-tpus

## AI 摘要

Google 团队在 MaxText 中用 Google Cloud TPU 从零复现了 AI2 的 OLMo 3 7B，覆盖 stage-1 预训练与 stage-2 mid-training，并在 held-out 指标上验证匹配。

## 正文

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.

MaxText and AI2 stage-2 training loss over the full anneal; the windowed gap closes to zero MaxText vs AI2 training CE across the full 47,684-step anneal (250-step means). The curves are indistinguishable from ~step 10k on. The lower panel shows the gap is carried almost entirely by the first ~8k steps, where the two shuffles have trained on mostly different data, then hugs zero: every 4k-window past 12k is ≤ +0.006, and the back third averages +0.0007, a tie. Overall mean Δ +0.0044.

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

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.

A subtler gotcha: the <|pad|> penalty

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 final landmark: held-out eval vs AI2

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:

Held-out C4 and perplexity for MaxText and AI2 at 24k/40k/47k, descending in lockstep at a flat ~+0.006 offset Held-out loss at three anneal landmarks. Left: C4 (solid) and 6-domain perplexity (dashed) fall in lockstep for both stacks. Right: the MaxText−AI2 gap stays tiny at every landmark, C4 flat at ~+0.006 nat and perplexity within +0.0027–0.0065. A faithful reproduction tracks the whole convergence curve, not just the endpoint.

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.)

Lessons for the next team

A multi-week, cross-stack reproduction teaches more than the final number. What we’d hand the next team:

Held-out eval at every landmark is non-negotiable. Training loss alone would have shipped a false “we beat the reference”: a data bug deflated it on repeated instances while held-out loss and accuracy never moved.

Keep optimizer state in fp32. weight_dtype=bfloat16 silently demoted Adam’s moments and added +0.93 loss over 1000 steps; the fp32 default closed it ~30×.

Resume must be exact. An off-by-one in resume-step detection desyncs data from parameters and is invisible on the loss curve; only a paired A/B diff catches it.

Decouple recipe from topology, and from hardware generation. JAX/XLA let the run ride a 512→128-device resize and repeated preemptions with no recipe change, then ran stage 2 with the same launcher on a different TPU generation (v5p, 57.4% MFU). Essential on a contended cluster.

Shape is a free lever. Reshaping heads (32×128 → 16×256) bought +12.4% throughput (44.2% → 49.6% MFU) at identical params and FLOPs, because head-dim 256 fills Ironwood’s 256×256 MXU. Worth checking before committing a long run.

Mind the multi-host footguns. jax.device_count() is local: thread a global TOTAL_DEVICES through or your effective batch silently shifts on a resize. And give auto-resume a real backoff (not a hard 60 s) to survive long scheduler queues.

What’s next

Stages 1 and 2 are done. Still ahead:

Stage 3: long-context adaptation, extending to 65k context with YaRN scaling on the global-attention layers.

Post-training via Tunix: SFT on Tulu-3, then GRPO for reasoning, reproducing OLMo-3-Instruct and OLMo-3-Think.

More architecture studies: beyond the head-dim result above, multi-token prediction (a small ~0.01 improvement in loss at λ=0.1, but −18% MFU from the extra remat) and compressed-expert “LatentMoE” layouts.

Reproduce it yourself

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_done

Stage 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_done

The 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.