MaxText 成功在 Google Cloud TPUs 上復現 Olmo 3 7B 預訓練

MaxText 的 stage-1 lm_loss 與 Ai2 已發布的 WandB 曲線相比，在 ~800k 步內的追蹤誤差在 ±0.012 以內。從 ~0.9M 步開始，MaxText 的曲線位於 Ai2 下方，且從未交叉回來。在 1.41M 步的 stage-1 損失比較中，差距在 ~800k 步以內維持在 ±0.012 以內，從 ~0.9M 步開始向負方向偏斜。在做 18k 步平滑之前，原始 2k 步區間的差距達到 −0.22。上述差距並未延伸到保留測試損失或下游準確率。

差距從哪裡來

在完整的 47,684 步中，與 Ai2 的訓練損失差距為 +0.0044，幾乎全部來自前 ~8k 步。文中指出，這段早期差距不是配方不匹配，而是兩邊的數據洗牌方式不同，早期模型訓練的主要是不同的數據。在 12.2M 實例的混合數據中，兩個隨機的 8k 步前綴只共享約 17% 的實例；在差距達到峰值的 1k 步，只共享約 2%。覆蓋範圍一旦重疊，差距就消失：第 12k 步之後，每個 4k 步窗口都落在 +0.0000 到 +0.006 之間，後三分之一的運行平均為 +0.0007，打成平手。

測量上也有一項不對稱：MaxText 的交叉熵計算會遮掉 <|pad|>，Ai2 則以接近零的成本把它算進去，這個偏差會把 Ai2 的曲線往下推，所以 +0.0044 比較像是同類比較下差距的上限。

評估設定

保留測試評估是在 16M tokens 的保留 C4-en 數據上做僅前向傳播評估，兩次運行使用相同批次。下游套件是八任務 lm-eval-harness，包含 MMLU、HellaSwag、ARC-easy/challenge、OpenBookQA、PIQA、BoolQ 和 WinoGrande。另有 Paloma 風格的掃描，涵蓋網頁、新聞、百科和混合領域，以及在相同輸入上計算的下一個 token 分佈距離（token-level KL）。

擴展與資源變動

JAX/XLA 堆疊有一項特性：配方與拓撲解耦。全局批次固定為 512 個實例、4.19M tokens/step，用多少晶片來承載則不固定。Ironwood 每個晶片有兩個 JAX 設備，因此 64 晶片的 4×4×4 切片會暴露 128 個設備。

在相同全局批次下，從 128 設備切片擴展到 512 設備（4 倍），1000 步測試中總吞吐量提高 3.99 倍，約等於 100% 的 strong scaling。SparseCore 卸載讓通訊不在關鍵路徑上。

在第 ~1.05M 步，容量損失四分之三，運行改在 128 設備切片（原來的四分之一大小）上恢復，全局批次與配方都沒有改。兩種配置下，每設備吞吐量都維持在 ~510–513 TFLOP/s/device，改變的只有每步的實際時間，從 0.76 秒變成 3.05 秒，與預期的 4 倍相符。

恢復的正確性

運行從第 19,500 步檢查點恢復，重新訓練了 127 個遺失的步驟，約相當於 47,684 步運行中六分鐘的重算量。重新訓練的損失與原始損失完全一致：全部 127 步的損失與困惑度 Δ 都是 0.000，而錯誤的恢復通常會產生約 0.01–0.4 的散布。無狀態恢復在 0/2 個工作程序時吻合，在 4 個工作程序時不吻合；Grain 自身的迭代器狀態檢查點，在每種工作程序數量下都是位元級精確。

後續階段與其他實驗

Stage 2（中訓練）沿用同一條管線，指向 Dolmino 索引，並從 warm-Adam 初始化開始。這個初始化由 stage-1 最終的全狀態檢查點建立，保留參數與 Adam mu/nu，並重置循環步驟與 LR 調度計數器。Dolmino 100B 索引則是從公開混合文件一次性建出。Stage 3 做長上下文適應，將上下文擴展到 65k，並在 global-attention 層使用 YaRN scaling。

Post-training 透過 Tunix 進行，先在 Tulu-3 上做 SFT，再做 GRPO 推理，用來復現 Olmo-3-Instruct 和 Olmo-3-Think。Multi-token prediction 在 λ=0.1 時帶來約 0.01 的損失改善，但因額外的 remat，MFU 下降 18%。文中也研究了包含 compressed-expert 的「LatentMoE」佈局。

比較工具（tools/wandb_csv_to_tensorboard.py、tools/compare_loss_curves.py、tools/eval_lm_loss.py）產生了本文中的所有表格。