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

在MaxText中重现Olmo 3 7B预训练：TPU大规模训练案例研究 - Google开发者博客

Google Tag Manager (noscript)

结束Google Tag Manager (noscript)

HTML

在MaxText中重现Olmo 3 7B预训练：TPU大规模训练案例研究

2026年9月24日

Gagik Amirkhanyan

高级软件工程师

Ran Ran

Aireen Mei

软件工程师

Matt Davidow

TPU推理软件工程团队

Google Cloud

Ai2团队

分享

• Facebook

• Twitter

• LinkedIn

• 邮件

由人工智能研究所（Ai2）开发的Olmo 3是一款最先进的全开放语言模型，采用现代架构和多阶段训练方案进行训练。为了评估MaxText在Google Cloud TPUs上的能力，我们的团队决定从零开始完全复现Ai2的Olmo 3 7B模型。我们选择Olmo 3是因为它结合了三个罕见共存的特性：它是一个在真实生产规模上训练的强大多模态7B模型，Ai2公开了完整的模型流程（包括数据、代码、配置、检查点、日志和评估），最后它还提供了独立的PyTorch和GPU参考基准，供我们测试MaxText和TPUs。

我们在Google Cloud TPUs的MaxText中成功复现了Ai2的Olmo 3 7B模型，包括第一阶段预训练和第二阶段中训练退火过程，并通过保留指标验证了与原始模型的匹配性（而不仅仅是损失曲线）：

主要亮点将在后续详细说明：

• PyTorch到JAX模型转换。Olmo 3的架构（重排序归一化块、QK归一化、3:1滑动/全局注意力）已移植到MaxText，并通过logit一致性检查验证：转换后的第0步检查点与HuggingFace参考在KL≈1.5e-3时匹配，达到"相同模型不同框架"的噪声基线；在8192-token上下文的bfloat16格式中，两者在top-1 token上98.75%时间达成一致。

• 能发现真实bug的验证机制。保留评估发现了一个数据加载器bug，使MaxText看起来似乎超越了参考模型；实际是记忆效应导致的性能提升。

• 多周运行的可靠性。检查点-恢复机制能完全复现运行：受控A/B测试显示恢复后每一步Δ=0.000；当主机故障导致第二阶段运行中途终止时，恢复运行重新训练了127步，日志中的损失和困惑度Δ保持0.000。

• 运行中动态调整训练任务规模。在约1.05M步时失去四分之三算力，运行在四分之一规模切片上恢复（相同脚本run_olmo3_7b_stage1.sh按设备批量大小调整保持GBS恒定），设备吞吐量保持在1%以内（≈100%强扩展性，双向测量）。

• 中途更换TPU代际的验证。第二阶段将相同启动器指向v5p而非Ironwood，仅改变设备类型，仍保持57.4% MFU。

• 自我投资的性能优化。通过SparseCore集合卸载、remat调优和最优分片，在Ironwood上实现7B规模44.5% MFU，大致相当于计算预算的三分之一被回收。

• TPU协同设计：在相同质量下更快。将注意力机制从32头×128头维度调整为16×256头维度（参数和FLOPs保持不变），由于256×256 MXU完全利用Ironwood的硬件特性，运行速度提升12.4%；其损失曲线在120B token（30k步）范围内与原始模型完全一致。这是个侧面消融实验；复现过程保留了原始架构。

从Ai2的step-0 PyTorch权重和相同的核配方出发，MaxText运行在完整约5.93万亿token/141万步预算范围内追踪了Ai2公布的损失曲线，并在stage-1结束时与之达成一致。我们甚至简化了两个配方细节（单个余弦LR调度而非Ai2的拼接方式，以及公开发布的数据混合；参见下方配方），匹配结果依然成立。本文其余内容将说明这些是如何构建、测量的，并在其中一个具有启发性的案例中几乎被伪造。

为何要复现Olmo 3？

Olmo 3是少数真正开放的前沿级语言模型之一：开放权重、开放数据，以及包含Weights & Biases上公开参考运行的完整训练配方。在保留指标而非仅损失曲线上独立复现该训练结果，是MaxText栈（优化器、损失函数、数据管道、数值计算）忠实性的重要证据，而不仅仅是“看起来在训练”。

MaxText是为TPU构建的JAX/XLA大语言模型训练框架。我们试图回答的问题是：能否在JAX-TPU上忠实复现PyTorch-GPU配方？能否在关键指标而非逐位对齐上达成匹配？又该如何证明这一点？

Olmo 3的训练配方是三阶段课程：通用预训练、中期训练（退火）和长上下文适配。本文涵盖第一阶段（约5.9万亿token预训练）和第二阶段（中期训练），这两阶段均端到端训练并匹配Ai2的参考结果。第三阶段及后训练（通过Tunix的SFT/RL）是已编写但尚未运行的配方。

Olmo-3预训练课程：本文复现第一阶段（Ironwood）和第二阶段退火（v5p）；第三阶段长上下文（seq 65k, YaRN）及后训练（SFT后GRPO via Tunix）将在后续文章中介绍。Olmo-3课程中，第一、二阶段在本文中复现；第三阶段及后训练将在后续文章中介绍。

配方

Olmo 3 7B是32层、4096维的密集变压器，包含几个非标准选择：一种“重新排序的归一化”模块、QK归一化，以及滑动窗口注意力与全局注意力3:1的混合。MaxText配置（olmo3-7b-pt.yml，用于第一、二阶段）与之完全匹配：

训练配方与Olmo-core的pretrain-1.py完全一致：曲线对齐所需匹配的参数设置：

我们从Ai2的step-0 PyTorch检查点开始训练，转换为Orbax格式，因此MaxText从与参考完全相同的权重开始训练。转换本身即为第一个检查点：在转换后的权重上进行前向传播，与HuggingFace参考在KL≈1.5e-3时达到9/10的top-10 token重叠，这构成了“同一模型、不同框架”的噪声基线。

匹配是否依赖于继承Ai2的初始化？显然不是。作为独立验证，我们还从MaxText自身的随机初始化训练了约5万步（占总训练时长的3.5%），其训练损失紧密跟踪Ai2公布的曲线，略低于该曲线。这是训练损失的抽样检查，而非完整复现，但表明匹配并不依赖于从Ai2权重开始。

数据管道与Olmo-core完全一致：对所有文档进行分词并连接（文档间插入EOS），切分为非重叠的8192-token实例，使用固定种子全局打乱索引，并应用n-gram重复过滤器，对包含超过32个重复n-gram的实例进行掩码。MaxText的dataset_type=olmo_grain（基于Grain实现）即实现了该流程。

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.

如何进行测量。所有评估均在对齐步骤的检查点对上运行：实时运行的 MaxText 与相同步骤的 Ai2 检查点（即其公开的 HuggingFace 版本 allenai/Olmo-3-1025-7B@stage1-step{N} 转换为 Orbax，转换过程可复现 PyTorch 参考值 KL ≤ 1.8e-3，即框架噪声基线）。lm-eval 使用标准 lm-eval-harness（5 个样本的 MMLU，其他使用默认值）；σ 是每个任务的 harness 标准差，通过四则运算组合用于计算差异。

在 stage-1 结束时，所有表面都一致认为两种配方是可互换的：

第四个表面，令牌级 KL，是唯一一个不是极小的数字：在相同输入上平均为 0.389 纳特，约为框架噪声基线的 200 倍。这符合同一配方的两次独立运行（总技能相同，但概率质量分配不同）的预期，这也是我们称之为“运行间噪声”而非“逐位一致”的原因；详细分析见 §G。

下游准确率差距在所有六个地标点上从未超过 ±0.005 宏观值，符号翻转四次，这正是两个忠实运行仅在 RNG 和数值上存在差异时预期的随机游走（每个地标点的表格见附录 D）：

8 任务宏观准确率，MaxText 与 Ai2 在六个地标点的对比，每个地标点的差异值。下游能力在每个地标点上都是可互换的。与训练损失不同，准确率从不单调发散；差异值在 ±0.005 范围内随机游走，最终结束于 +0.0002。

看起来像胜利的错误

事情变得有趣的地方在这里。从约 0.9M 步开始，MaxText 的训练损失逐渐低于 Ai2 的，并且再也没有回升；在约 1.25M 步之后，它明显低于 Ai2，平均低 0.06，某些几百步的区间甚至低达 0.25。仅观察训练损失叠加图，你会得出结论：MaxText 已经领先了。

实际上并非如此。在括号检查点的保留 C4 损失是平局的（1,000k 时 Δ −0.004，stage-1 结束时 +0.003），而在 1,350k 处下游准确率略微偏向 Ai2。训练损失在下降，而泛化能力没有变化。这是记忆化的特征：模型看到了某些序列不止一次，并在重复时获得了低损失。

顶部：MaxText 的训练损失降至 1.63，而 Ai2 保持平稳。底部：保留的 C4 差值在整个过程中保持在零附近。整个故事在一个图中。顶部：在 1.24M–1.41M 的窗口中，MaxText 的训练损失在 2k 步均值上反复降至 1.63（Δ −0.22；在数百步区间内低至 −0.25），而 Ai2 的保持平稳。看起来像是彻底的胜利。底部：训练损失 Δ（蓝色）逐渐变负，但保留的 C4 损失（绿色菱形）从未超出 ±0.02 范围。训练损失下降了；泛化能力没有变化。

原因是 Grain 数据加载器中的双重分片错误。MaxText 的 Olmo 加载器将 ShardOptions(shard_index, shard_count) 传递给了 Grain DataLoader，而索引采样器已经在内部进行分片。Grain 的 shard_options 不仅记录元数据，还会重新调整采样器的索引流。当 shard_count=32 时，数据光标以 32 倍的速度前进，导致 stage-1 不再是一个干净的 epoch，而是变成了 Poisson(≈1) 的有放回重采样：大约 37% 的语料库从未被看到，37% 被看到一次，26% 被看到两次或更多次。令牌预算保持不变（约 5.9T 实际令牌），这就是为什么损失仍然与 Ai2 全局一致，但重复的实例恰好在它们重复出现的地方降低了训练损失。

由此得出的两个教训是：

• 训练损失并不能证明收敛。我们之所以没有发布“MaxText优于基准”的错误声明，唯一原因是我们在每个关键节点都坚持使用保留集进行评估。在保留集C4和所有8个下游任务中，记忆化下降现象都不可见。

• 真实复现需要预留"错误预算"。修复方法（grain.sharding.NoSharding()，让采样器接管所有分片）仅需一行代码。找到这个错误需要A/B测试框架、一个能复现分片数>1时发散现象的单元测试，以及硬件重跑验证。

在验证修复过程中，我们发现了第二个独立错误：恢复步骤检测存在越界错误。检查点目录编号为N，但训练循环在第N次迭代完成后才会写入目录N，导致模型恢复到N+1步时，数据加载器却从第N批开始恢复，重新训练一个批次后，始终落后一步。应用两个修复后，检查点-恢复运行能完美复现无中断运行：所有99个步骤的记录损失Δ值均为0.000。（该A/B测试使用单工作者数据加载；多工作者场景在阶段2出现，我们已在此阶段解决，详见阶段2。）两个错误都有回归测试用例，旧代码会失败，修复后会通过。

我们让正在进行的阶段1运行按原样完成：它已完成85%，修复无法处理已读取的数据，重新启动将损失约120万步的计算量。上述验证表明该错误对可观察精度没有影响；修复仅针对未来运行。

在Ironwood上的性能和扩展性

复现数学公式只是任务的一半；另一半是让其运行得足够快，并在集群变化时保持速度。在140万步运行期间，作业被抢占、重新调度和调整大小多次，整个系统栈必须在不修改训练方案的前提下吸收所有变化。

提取MFU性能

在Ironwood上使用7B参数模型、每设备批量4时，原始架构（"variant D"，见附录J的消融实验）达到44.5%的MFU（510-513 TFLOP/s/设备）。仅改变头部维度的形状参数即可提升至49%；详见下文"头部维度"。影响性能提升的主要因素依次为：

• Ironwood XLA标志 + SparseCore卸载：将集合操作（all-gather、2D all-gather、reduce-scatter）卸载到SparseCore，配合v7x专用XLA标志，使MFU从41%提升至44.5%，且损失保持不变。（完整标志列表见附录H。）

• 扩展重计算：检查点保存注意力和MLP投影（qkv_proj、q/k/v_proj、out_proj、mlpwi_0、mlpwo、context）使激活内存预算达标；再添加mlpwi_1时HBM溢出18GB。

• Splash注意力 + Tokamax（2048-token块）。

• 此规模下分片轴无关紧要：纯FSDP、4-FSDP×32-DP和8-FSDP×16-DP在128设备上均达到约1.5 TFLOP/s；纯FSDP因简单性更优。芯片内张量并行（TP=2）带来净损失：半批量时MFU下降1.6%，全批量时OOM（FSDP=64 per-chip权重状态）。

扩展性及为何是免费的

JAX/XLA栈在此最宝贵的特性是训练方案与拓扑结构解耦。全局批量（512实例，4.19M tokens/step）固定，但分配的芯片数量可变。（单位说明：Ironwood每芯片包含两个JAX设备，因此64芯片的4×4×4切片暴露128设备；我们同时引用两种表述。）

• 扩展验证：在相同全局批量大小下，从128设备切片扩展到512设备切片（4×）时，总吞吐量提升了3.99倍，1000步测试中实现了约100%的强扩展性。强扩展性是更具挑战的方向：每个设备现在每步只执行四分之一的工作量，而集体通信操作需要跨越4×数量的设备，因此可用于重叠通信的计算资源更少。SparseCore卸载技术仍然确保了通信操作不阻塞关键路径。

• 生产环境中的收缩：在约1.05M步时，我们失去了四分之三的算力，运行自动切换到128设备切片（原规模的四分之一），在相同全局批量大小下继续执行，且未更改任何配置。每设备吞吐量保持不变（在两种情况下均为约510–513 TFLOP/s/设备）；仅每步的墙钟时间发生变化（0.76秒 → 3.05秒，符合预期的4×变化）。

左图：总吞吐量从128台设备扩展到512台设备时提升了3.99倍。右图：每台设备的TFLOP/s在调整规模后保持不变。

这正是长运行能够在竞争性集群中持续执行的关键：利用任何可用的算力，保持数学计算完全一致。第二阶段将这一理念扩展到TPU代际之间（详见下文）。

自动恢复：跨越数周的运行

一个1.4M步的运行可能会被中断。我们通过带有 resume_until_done 循环的机制驱动它，该机制在抢占发生时自动重新提交任务，并从最新的Orbax检查点恢复执行：

• 每2000步进行一次检查点保存，因此抢占最多导致约2000步的重新计算（在大切片上仅需数分钟）。为重试循环设置真正的退避机制：早期版本在Kueue压力控制下耗尽MAX_RETRIES=50次重试；可配置的RETRY_BACKOFF_SECONDS（默认300秒）使其能够跨越数小时的调度空窗期。

• GCS上的TensorBoard是事实来源：kubectl logs只能查看当前Pod的历史记录。本文所有表格均来自GCS持久化存储的TB事件，而非实时日志。

• 恢复必须完全精确，否则会静默破坏运行。读取错误数据或从错误步数重启的恢复在损失曲线上看似正常，但实际运行已发生改变；这正是上文提到的两个数据加载器bug。修复后恢复过程完全精确，下文的A/B对照实验就是证明。

有bug的恢复会导致连续运行出现偏差；修复后的恢复在每一步都精确为零。在128设备（64芯片）上的受控A/B测试中：从第100步的检查点恢复，与未中断的运行进行对比。偏移一步的bug（红色）导致数据与参数出现偏差（平均|Δ| 0.044，最大0.22），且永远无法重新收敛。修复后（绿色）在所有99步中均为0.000。恢复bug在常规损失曲线上是不可见的；只有通过这种配对差异对比才能发现。

由于Ironwood的矩阵乘法单元（MXU）是一个256x256的系综阵列，标准的head-dim为128会导致阵列的一半在注意力QK矩阵乘法过程中处于空闲状态。将head-dim重塑为256可使张量维度与硬件完全对齐，彻底消除空闲计算周期。这在保持参数和FLOPs完全一致的情况下，带来了+12.4%的吞吐量提升（571 vs 508 TFLOP/s/设备，或49.6% vs 44.2% MFU）。这是一个无需额外成本的加速提升，在投入长期运行前实施该优化具有极高价值。（损失曲线和种子注意事项详见附录L。）

值得单独成段的注意事项

优化器数据类型是静默且代价高昂的。设置weight_dtype=bfloat16会通过mu_dtype继承机制静默地将Adam的m/v矩量降级，导致1000步后损失增加+0.93：bf16的约3位有效数字会丢失每个微小预热更新的一小部分，这种影响会不断累积。将weight_dtype保留为默认的float32可将差距缩小30倍。这是项目中最大的“为何不匹配”时刻。

计算

阶段1的计算消耗了约77,000 Ironwood芯片小时的步骤计算（约3,200芯片日），再加上检查点、评估和重启阶段的开销。芯片小时是固定不变的单位：每个步骤的芯片秒数与切片无关（0.76秒×256芯片≈3.05秒×64芯片≈195芯片秒），而墙钟时间取决于切片。主体计算在4×8×8切片（256芯片/512设备）上运行，77,000芯片小时相当于约12.5天；加上预emption后在64芯片切片上的运行时间和排队时间，实际日历时间延长至数周。在预运行30% MFU预算下，相同token量需要约50%更多的芯片时间（相同步骤时间核算下约113,000芯片小时；附录I的6·N·D规划行显示约100,000），因此性能优化工作挽回了约三分之一的资源消耗。阶段2相对便宜：约5,000 v5p芯片小时（在128芯片v5p-256上约39小时；详见阶段2和附录I）。

阶段2：中训练（退火）

在阶段1匹配后，我们进入阶段2：中训练，即预热-稳定-衰减（WSD）计划的最终衰减阶段。阶段1模型在Dolmino 100B混合数据集（高质量数学、代码、推理和精选网页）上进行退火，同时学习率从2.0712e-4线性衰减至0。这是Olmo-3的高质量数据转化为能力提升的关键阶段，因此忠实的堆栈配置也必须匹配。我们已通过Ai2的中训练参考配置（运行zxv811e1，由Olmo-core的Olmo-3-1025-7B-midtrain.py生成）验证了该方案；所有超参数完全匹配：

Warm-Adam初始化。Ai2从阶段1的最终3e-5立即重新预热到2.0712e-4，无预热且load_optim_state=True；加载的Adam二阶矩量合理地吸收了这一跳跃。我们通过一次性检查点手术（保留参数+mu/nu，将循环步数和学习率计划计数器置零）复现了这一过程，因此运行从峰值重新启动计划，匹配Ai2的初始化配置。第0步损失约为1.53，而非冷启动尖峰。

不同TPU代际，相同启动器：v5p上57.4% MFU

Stage 2也迁移了硬件世代。由于Ironwood的算力被分配到其他任务，我们使用相同的启动脚本、针对Ironwood优化的XLA标志等，将这些配置指向TPU v5p切片（v5p-256，128芯片），仅修改了XPK设备类型。在未进行v5p特定优化的情况下，达到了57.4%的MFU（v5p每芯片263 TFLOP/s的中位数，占其459峰值的57.4%），高于阶段1在Ironwood上44.5%的表现。这是因为7B模型更容易饱和较旧芯片，而v5p的峰值FLOPs是其5倍。并且这一表现保持稳定：整个运行中，从第25百分位到第90百分位，每芯片吞吐量在263.0至263.9 TFLOP/s之间波动，47,684步的差异仅为0.4%。结合阶段1的调整，这完整体现了可移植性：该方案与切片拓扑和TPU世代完全解耦。你可以在任何可用的算力上运行。

在整个47,684步训练中，与Ai2相比，整体训练损失差距为+0.0044，几乎完全由前8,000步造成。这一早期差距并非方案不匹配：初期时，两种数据洗牌主要训练在不同数据上。12.2M实例混合数据的两个随机8,000步前缀仅共享约17%的实例（在1,000步时差距最大，仅共享约2%）。一旦覆盖重叠，差距逐渐消失：从第12,000步后，每个4,000步窗口的差距在+0.0000至+0.006之间，后三分之一阶段平均差距为+0.0007，达到平局。需要特别注意的一个度量不对称性：我们的CE掩码处理<|pad|>（下一节详述），而Ai2的处理成本接近于零，这种偏差会将Ai2的曲线向下推移，因此+0.0044实际上是同类对比差距的上限。

MaxText和Ai2阶段2训练损失在整个退火过程中的表现；窗口化差距在MaxText与Ai2训练CE的完整47,684步退火过程中（250步均值）缩小至零。从约第10,000步开始，两条曲线几乎无法区分。下图显示差距几乎完全由前8,000步造成，此时两种数据洗牌主要训练在不同数据上，随后差距归零：从第12,000步后，每个4,000步窗口的差距≤+0.006，后三分之一阶段平均差距为+0.0007，达到平局。整体平均差距Δ为+0.0044。

仅阶段2的恢复漏洞：也要检查数据迭代器

阶段2暴露了一个阶段1修复未覆盖的恢复漏洞：阶段1的修复仅在grain_worker_count=1时保证恢复精确性。阶段2运行4个工作者（与Ai2匹配），此时无状态恢复（从步骤重新推导数据偏移）会发散，因为新加载器从偏移开始时，会以不同于不间断流的方式交错工作者记录。一个奇偶性测试验证了这一点：无状态恢复在0/2工作者时匹配，在4工作者时不匹配；Grain自身的迭代器状态检查点在每个工作者计数时都是位精确的。

因此在阶段2中，我们通过Grain的GrainCheckpointHandler将olmo_grain连接起来：现在检查点除了模型项外，还会携带一个迭代器项，恢复时会原子性地同时恢复权重、优化器状态和精确的数据迭代器位置。在竞争性集群中，当47,684步运行会被多次抢占时，这正是确保完整运行为一个干净的完整周期而非嘈杂的有放回重采样的关键。

证明意外到来。在运行了16小时后，主机故障导致作业集在第19,627步终止（直至故障前损失曲线始终平滑：基础设施短暂波动，而非模型发散）。运行从第19,500步检查点恢复，重新训练了127个丢失的步骤，整个47,684步训练中仅耗时约6分钟。由于这些步骤在崩溃前已记录在日志中，抢占机制为我们提供了免费的配对差异，重新训练后的损失与原始结果完全一致：在127个步骤中，损失和困惑度的Δ均为0.000，而中断恢复通常会产生约0.01–0.4的波动。这一精确重放发生在grain_worker_count=4的配置下，这正是无状态恢复会失败的场景。中断运行能够完全重放并非偶然；这是三个确定性机制共同作用的结果：XLA编译的TPU程序、Grain迭代器状态和Orbax恢复。

更微妙的陷阱：<|pad|>惩罚

早期有一个数值看起来异常：在没有填充掩码的情况下，我们第二阶段训练损失在olmOCR PDF分片上飙升至约7.5（<|pad|>，token 100277，约占中期训练token的2%），而Ai2的曲线则平滑（最大约1.96）。原因在上游：我们第一阶段的n-gram重复过滤器屏蔽了所有填充窗口，因此模型从未学习预测<|pad|>，在强制预测时损失约17，而Ai2的损失则几乎不受填充惩罚影响。我们对填充不敏感模型的忠实选择是：在损失中屏蔽填充token（olmo_pad_token_id=100277）；在相同批次上进行配对LR=0的评估确认模型其余部分等效（Δ −0.001）。“匹配参考配方”有时意味着要考虑上游阶段的差异，而非盲目复制当前阶段配置。

最终里程碑：保留集评估 vs Ai2

完整的100B运行在第47,683步完成，学习率衰减至0，期间自动恢复了所有中断。真正重要的问题不是训练曲线，而是最终模型是否在从未训练过的数据上与Ai2的中期训练检查点匹配。

我们使用统一的评估栈，将最终检查点与Ai2的Olmo-3-1025-7B第二阶段第47684步进行对比：保留集C4损失、多领域困惑度（Paloma风格）和8任务lm-eval套件。

损失侧指标在稳定值~+0.006 nat上达成一致：训练CE、保留集C4和困惑度均值出现相同的微小偏移。这种一致性指向跨框架数据排序问题，而非bug；真实bug通常会导致单个领域异常，而非所有领域同等程度异常。按领域划分的困惑度验证了这一点：wikitext103的偏移为−0.0005，最宽领域是wikipedia，偏移为+0.0135（附录M）。

关键的是，这种损失偏移并未转化为能力差距。下游准确率持平：宏平均−0.0023，8个任务在评估噪声范围内双向波动（每任务指标：lm-eval默认，定义acc_norm时使用，否则使用acc）：

这不仅是终点线。在退火过程中三个里程碑处评估两个模型栈，显示保留集差距始终平坦且微小：C4在所有三个里程碑处保持~+0.006 nat，6领域困惑度维持在+0.0027–0.0065范围内，两个模型随着学习率衰减至0同步提升：

对下一团队的启示

跨栈多周复现不仅教会了最终结果，更揭示了关键经验。我们希望传递给下一团队的要点：

• 每个关键点的保留评估是不可妥协的。仅依赖训练损失会误导得出“我们超越了基准”的错误结论：数据缺陷在重复实例中会压缩训练损失，但保留损失和准确率始终稳定。

• 保持优化器状态为fp32。将weight_dtype设置为bfloat16会静默降低Adam的矩估计，导致1000步增加+0.93损失；而fp32默认值可将差距缩小约30倍。

• 恢复必须精确。恢复步骤检测中1个单位的偏差会导致数据与参数脱节，这种影响在损失曲线上完全不可见，只有通过A/B配对差异分析才能发现。

• 将配方与拓扑结构、硬件生成解耦。JAX/XLA使运行可在512→128设备的规模调整和重复抢占中保持配方不变，随后使用相同启动器在不同TPU代际(v5p, 57.4% MFU)上运行第二阶段。在竞争集群中这是关键要素。

• 形状是免费的杠杆。重塑头结构(32×128 → 16×256)在相同参数和FLOPs下提升+12.4%吞吐量(44.2% → 49.6% MFU)，因为256头维度完全填充了Ironwood的256×256 MXU。在长期运行前值得验证。

• 注意多主机陷阱。jax.device_count()是本地值：需要全局传递TOTAL_DEVICES变量，否则在规模调整时有效批量会静默偏移。为自动恢复设置真实退避机制(而非硬性60秒)可应对长队列调度。

后续计划

第一、二阶段已完成。仍需推进：

• 第三阶段：长上下文适配，通过全局注意力层的YaRN扩展将上下文延伸至65k。

• 通过Tunix进行微调：在Tulu-3上进行SFT，然后使用GRPO进行推理训练，复现Olmo-3-Instruct和Olmo-3-Think。

• 更多架构研究：除上述头维度结果外，还包括多标记预测(λ=0.1时损失小幅度改善~0.01，但因额外remat导致MFU下降18%)和压缩专家"LatentMoE"布局。

亲自复现

配方、启动器和转换工具均集成在MaxText中。运行结构如下：

$ maxtext run --recipe <recipe-name> --tpu <tpu-name> --batch-size 1024

# 将 Ai2 的 step-0 权重转换为 Orbax，然后：
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
# 驱动完整训练，自动重新提交被抢占任务：
STEPS_OVERRIDE=1414078 bash src/maxtext/trainers/pre_train/scripts/olmo/xpk_olmo3_7b_stage1.sh resume_until_done

Shell

已复制

阶段2（中训练）复用相同流水线，指向 Dolmino 索引，从 warm-Adam 初始化开始：

# 1. 从 stage-1 最终完整状态检查点构建 warm-Adam stage-2 初始化
#    （保留参数 + Adam mu/nu，重置循环步数 + 学习率调度计数器）：
python3 -m maxtext.utils.olmo3_build_stage2_init          # 脚本头部包含 src/dst 路径
# 2. （仅一次）从公共混合文件构建 Dolmino 100B 索引：
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. 启动完整的 100B 烧入训练。启动器默认使用 Ai2-faithful
#    方案（学习率 2.0712e-4 → 0 线性，预热 0，47,684 步，批量 256，种子 1337，
#    填充掩码 100277），并精确恢复（权重 + 优化器 + 数据迭代器）：
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

启动器会固定配方表中的所有超参数，导出 Ironwood XLA 标志集，并根据全局设备数量计算每设备批量。比较工具（tools/wandb_csv_to_tensorboard.py, tools/compare_loss_curves.py, tools/eval_lm_loss.py）生成本文所有表格：指向两个 TensorBoard 目录即可打印差异值。

核心结论：基于 GPU 的 PyTorch 预训练方案在 TPU 的 JAX 中能忠实复现，但只有在测量泛化能力而非仅训练损失时才称得上“忠实”。我们最重要的方法论选择是每个关键节点都进行保留数据集的评估验证。这发现了伪装成胜利的数据错误，并让我们能正大光明地说“已复现”。

发布于：

• AI

• 云服务

• 案例研究

• 操作指南

• 最佳实践

• 学习

• 探索

• 影响力

上一篇

下一篇

相关文章

列表

AI

云服务

公告

解决问题

使用 Google Cloud API Gateway 将 REST API 转换为 MCP 工具

案例研究

最佳实践

最强 AI Agents 挑战赛提交背后的4个工程模式

2026年9月2日

教程

在 TPUs 上加速视频扩散的时空注意力

2026年9月30日

导航点