Sept. 24, 2026 (英语).
Olmo 3，由Allen Institute for AI(Ai2)开发，是一种最先进，完全开放的语言模型，经过现代建筑和多阶段培训配方的训练。 为了评估MaxText在Google Cloud TPU上的能力，我们选择了奥尔莫3，因为它结合了很少同时出现的三个属性。 这是一个强大的现代7B模型，经过实际生产规模的培训。 Ai2披露了完整的模型流量，包括数据，代码，配置，检查站，日志，以及评价。 最后，它给了我们一个独立的PyTorch和GPU的参考，我们可以用它来测试MaxText和TPU.
我们重现了Ai2的 Olmo 3 7B 在 MaxText 上的 Google Cloud TPU，无论是第1阶段前训练还是第2阶段中训练的肛门
主要要点，后面各有详细论述：
- PyTorch JAX型号转换。 Olmo 3的架构(重新排序-规范块，QK-规范，3:1滑动/全局关注)移植到MaxText，并通过对数-均等检查进行验证：转换后的台阶-0检查点与KL QHuggingFace参考点吻合 0.5e-3,"相同模型，不同框架"噪声层，在bfloat16的完整8192-token上下文中，两人同意时间最高值-1的98.75%.
- 能够捕捉到真虫的核查 Hold-out evals 捕捉到一个数据加载器错误，使得MaxText看起来像是在击败引用；增益是回忆。
- 多星期的可靠性检查点和重置精确地重置了运行：在恢复后的每一步都有一个控制的A/B 显示 = 0.000，当一个宿主故障杀死了阶段-2的运行中程时，恢复运行在XQ = 0.000的记录损失和perplexity重新训练了127步。
- 调整中途的训练任务。 在步骤~1.05M中，我们失去了四分之三的容量；运行恢复在四分之一的片段上，没有改变食谱的大小(相同的脚本运行 olmo3 7b stage1.sh比例尺每个设备批量大小保持GBS恒定)，每个设备的吞吐量保留在1%之内(X%100%强缩放，双向测量)。
- 改变 TPU 生成的中继器。 第2阶段将相同的发射装置指向v5p而不是Ironwood，只改变设备类型，维持了57.4%的MFU.
- 自己付的功劳在7B的铁木上，44.5%的MFU通过SparseCore集体卸载，重制调和，以及最优的硬化，大约三分之一的计算预算回购。
- 为TPU共同设计：在同样质量下更快。 将注意力从32头×头-dim 128重排到16×256重排，在相同的参数和FLOP重排时，运行速度加快了+12.4%，因为头-dim 256充分利用了Ironwood的256×256MXU，其损失曲线通过120B令牌(30k步)与原始的相匹配。 这是一种侧式建筑；复制保留了原始建筑。
从Ai2的步-0 PyTorch重量和同一核心食谱开始，MaxText运行了Ai2公布的亏损曲线，覆盖了完整的~5.93T-token / 1.41M步预算，并在第1阶段末端的顶部运行获得。 我们甚至简化了两个食谱细节(一个Ai2缝合两个的单一大同LR时间表，以及公开的发布数据组合；参见下面的食谱)，以及无论如何举行的比赛。 其余部分是这些东西是如何建造、测量的，在一个有启发性的例子中，几乎是伪造的。
Olmo 3是少数真正开放的前沿类语言模型之一：开放权重，开放数据，以及一个在Weights & Biases上公开查询的完全指定的训练食谱。 匹配独立训练的运行，而不是仅仅按照损失曲线进行，这有力地证明了MaxText堆(优化、损失、数据管道、数字)是忠实的，而不仅仅是“似乎是训练 ”。
MaxText是为TPU建造的JAX/XLA LLM训练框架。 我们所要回答的问题：PyTorch-on-GPU的食谱能否在JAX-on-TPU中忠实地复制，在重要的而不是比特比特比特的衡量标准上匹配，你如何证明？
Olmo 3的食谱是三阶段课程：一般的预训，中训(ananaling)，以及长文适应。 这个职位涵盖第1阶段(~5.9T-token预训跑)和第2阶段(中训)，既经过训练结束，又与Ai2的参考文献相匹配。 以及训练后(SFT/RL通过Tunix)
Olmo-3 预训课程：第1阶段(Ironwood)和第2阶段(v5p)在本文章中转载；第3阶段长文(seq 65k,YARN)和后训练(SFT之后通过Tunix的GRPO)接下来
奥尔莫-3课程。 第一阶段和第二阶段在此员额中复制；第三阶段和后培训是下一个阶段。
Olmo 3 7B是一个32层，4096深密的变压器，有一些非标准的选择：一个“重排规范”块，QK-Norm，以及3:1滑窗与全球关注的组合。 MaxText配置(olmo3-7b-pt.yml，用于第1阶段和第2阶段)与其完全吻合：
训练食谱与Olmo-core的Pretrain-1.py相仿；
我们从Ai2的步数-0PyTorch检查站开始训练，转换为Orbax，因此MaxText从与参考文献完全相同的重量开始。 转换本身是第一个检查点：转换后的重量的前传与KL ≈ 1.5e-3的HuggingFace参考值相匹配，9/10上-10符号重叠，即“相同的模型，不同的框架”噪音层。
比赛是否取决于继承Ai2的初始化？显然没有。 作为独立检查，我们还从MaxText自己的随机输入器中为~50k步(占地平线的3.5%)训练了一次运行；其训练损失跟踪了Ai2的已公布曲线，在它下面运行一触。 这是一次训练损失现场检查
数据管道精确地镜像了Olmo-core：将所有文件(EOS)标注在两者之间，切成非重叠的8192个托肯实例，用固定的种子在全球对索引进行洗涤，并应用一个n克的重复过滤器，用>32的重复n克来掩盖实例。 MaxText 的数据集 type=olmo grain( 在谷物上建造) 执行此功能。
与Ai2运行的两个故意的分歧 (1)LR时刻表：Ai2最初计划~5T令牌，并将飞行中程扩展至~5.93T的终极视野，因此其LR的微量线缝合了两条余弦曲线(可见于其公开的WandB运行)；我们运行了一条余弦线全程。 (2)数据：我们用公开的发布 Olmo-3混音器来训练，它省略了Ai2内部运行所看到的一小部分(小于象征性预算的0.5%，大多为发布文件列表中缺失的s2pdfsshards). 两者都是我们选择的简化，而不是意外，而MaxText仍然在每个被搁置的表面匹配。 这也是我们为什么说“重复到运行中的噪声，”而不是位元(见§G中的KL分析)。
Olmo 3在我们开始的时候不是在 MaxText中；复制增加了，并且上游，下面的一切。 本文章结尾的"自行转载"是"配置"，而不是代码。
- 模型本身：重排规范块，QK规范，以及3:1滑动/全球关注模式(#3004,#3112).
- 匹配 Olmo-core 语义的跳跃优化器，降为 Bessel 修正的 running std (在 128 步窗口上滑动 6/ 5) (# 3490)。
- Z - 损失(# 3211)和每参数的重量级掩体，因此可以排除嵌入(# 3280).
- olmo grain数据管道(# 3749)：随机访问读取了预托的硬盘，一个种子全球指数洗牌，上面有指纹保护符，防止重启时的无声数据交换，n格重复过滤器，以及(从第2阶段开始)Grain interator-State检查。
- HF Orbax检查站转换，并配有对数-对数检查，这是本站每个"框架噪音地板"号码(#3112,#3832)背后的工具。
- 舞台-1和2发射器(env驱动运行的脚本+XPK包装器，并提交/监视器/恢复直到 done)(#3886).
- TensorBoard 等值标记(优化/步骤 skipped,perf/total totokens)，因此Ai2的W&B仪表板上的每一度量都有一个MaxText对应比对。
头条是单倍数：MaxText的阶段-1 lm loss vs Ai2 已公布的 WandB 曲线，步对齐并被宾入2k步表示。 通过~800k在±0.012内步入双轨；从~0.9M MaxText在Ai2下方的边缘并永不交叉，数据错误的第一个标志在下一节解析。
MaxText vs Ai2 阶段-1在1.41M 步骤中丢失，下面板有 18k 步的偏差
顶部：曲线在这个尺度上是无法区分的，直到尾部。 底部：空隙停留在±0.012到~800k内，由于数据错误重复堆积而从~0.9M处倾斜为负，潜过1.25M，在18k步平滑两面板绘制前在原始的2k步槽中达到−0.22. 如下文各节所示，所有这一切都没有达到搁置损失或下游准确度。 (Per-landmark表：附录A；全曲线承诺为olmo stage1 loss curve.tsv.).
但光是损失曲线就是一个薄弱的证明：在训练损失上两跑就能匹配，因此我们在6个跨度915公里的地标上验证了四个独立表面的趋同：
- Hold-out C4 lm loss：在16M的代号C4-en上只转发eval，相同的批次都运行。
- 8任务lm-eval-harness套房：MMLU,HelaSwag,ARC-easy/challenge,OpenBookQA,PIQA,BoulQ,WinoGrande.
- 多域屏蔽perplexity：一个帕洛马式的横扫网络、新闻、百科全书和混合域。
- Token-level KL：在相同输入上的下一个token分布距离。
我们如何衡量。 所有evals运行在步配对的检查站上：直播的MaxText运行对同一步骤的Ai2检查点，即其公开的HuggingFace修订Allenai/Olmo-3-1025-7B@stage1-step{N}转换为Orbax(转换复制了PyTorch参考KL ≤ 1.8e-3，框架噪音层). lm-eval使用标准lm-eval-harness(5-shot MMLU，在别处默认); ç 是一个每个任务带 Stderr，组合为三角洲的四元。
在第一阶段结束时，每个表面都同意两种食谱可以互换：
第四层，象征级的KL，是一个不小的数字：相同的输入值为0. 389纳特，这也是我们为什么说“运行到运行的噪音 ”，而不是位元；分解在§G中。
下游准确度差距从未超过±0.005宏观，标志翻转四次
8 任务宏精确度，MaxText vs Ai2 跨越六个地标，每块地标三角洲
下游的能力在每个地标上都可以互换。 与训练损失不同，精度从不单调；三角洲随机行走在±0.005内部，以+0.0002结束。
这是有趣的。 从~0.9M步的MaxText的训练损失在Ai2下方倾斜，再也没有回转过；过去~1.25M路的训练损失明显下降，平均为−0.06路，几百步的训练损失高达−0.25路。 仅观察训练损失的覆辙
没有。 括号内检查站的C4损失被扣(++0.004,1000k,++0.03,1阶段末)，下游准确度为1 350k，稍为偏好。 训练损失在下降，这便是记忆的标志：模型看到一些序列不止一次
顶部：最大Text训练损失潜至1.63 而Ai2保持平稳。 底部：C4三角洲被搁置，整个期间接近零
整个故事在一个图。 顶部：在1.24M–1.41M窗口中，MaxText的训练损失在2k步(QQ−0.22;~0.25步(百步拉伸))上反复跌至1.63，而Ai2则保持平稳。 看起来像一个逃跑的胜利。 底部：训练-损失 QQ(蓝色)漂移负数，但被搁置 C4 l
本文由 AI图鉴自动抓取并翻译，版权归原作者所有；如需查看原文请点下方链接。
SEPT. 24, 2026
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 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.
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.
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 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.
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.
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 l
Auto-fetched and translated by AI Atlas. All rights belong to the original author.