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

MaxText 團隊利用 JAX 與 XLA 技術，在 Google Cloud TPUs 環境中從頭成功復現了由 Ai2 開發的 Olmo 3 7B 語言模型。這項實作嚴格對齊了原始的 PyTorch-on-GPU 參考實現，在預訓練與中訓練這兩個關鍵階段，所有保留驗證集（held-out evaluations）上的表現數據均精確匹配原始結果，證明了訓練精度上的高度一致性。

在效能方面，該實作達成了高達 57.4% 的模型浮點運算利用率（MFU）。系統在訓練過程中經歷了中途的集群調整（mid-run cluster resizes）以及跨代 TPU 遷移，在此類基礎設施變動下，系統依然穩定運作，且完全無需對訓練配方（recipe）進行任何修改或調整，展現了極高的架構韌性與硬體資源的有效運用。

此次實驗在工程實踐上驗證了全面保留驗證集的必要性。透過嚴格的驗證流程，團隊成功揪出一個靜默的資料載入器記憶錯誤（data-loader memorization bug）。這個錯誤會導致訓練損失數值虛低，若未被發現將掩蓋模型真實效能並造成性能提升的假象。這起錯誤的發現，凸顯了在大型語言模型訓練中，對訓練過程數據完整性進行嚴格監控的關鍵價值。