Google、MaxTextとTPUでOlmo 3 7Bの事前学習を再現した事例を公開

再現ではAi2のstep-0 PyTorchチェックポイントをOrbaxに変換し、同じ重みから学習を開始した。変換後のロジットはHuggingFaceの参照とKL ≈ 1.5e-3で一致した。stage-1は約5.93Tトークン、1.41Mステップに及び、ホールドアウトC4損失の差は+0.003、8タスク精度の差は+0.0002だった。stage-2は100Bトークンで、損失差は+0.006、精度差は−0.0023でいずれも誤差範囲とされている。

検証の中心はホールドアウト評価である。約0.9Mステップ以降、MaxTextの学習損失がAi2を下回り続けて優位に見えたが、ホールドアウトC4損失は同等のままだった。原因はGrainデータローダーの二重シャーディングで、データカーソルが32倍速く進み、コーパスの約37%が未使用、26%が2回以上使われていた。修正は `grain.sharding.NoSharding()` の1行である。再開ステップの検出にもoff-by-oneのバグがあり、両方の修正後はチェックポイントからの再開で99ステップすべてΔ = 0.000となった。

運用面では、約1.05Mステップで容量の4分の3を失い、4分の1サイズのスライスへレシピを変えずに再開した。デバイスあたりのスループットは1%以内で維持された。stage-2は同じランチャーをIronwoodからTPU v5pに向けただけで動作した。SparseCoreへのコレクティブオフロード、rematerializationの調整、最適なシャーディングでIronwoodのMFUは44.5%に達し、計算予算の約3分の1を取り戻したと述べている。

ハードウェアとの協調設計として、アテンションを32ヘッド×head-dim 128から16×256に変える実験も行った。パラメータ数とFLOPsは同一で、IronwoodのMXUが256×256であるため12.4%高速になり、損失曲線は120Bトークンまで元と一致した。ただしこれは副次的なアブレーションで、再現本体は元のアーキテクチャを使っている。

今後はYaRNによる65kコンテキストへのstage 3と、Tunixを使ったSFTおよびGRPOによるポストトレーニングが予定されている。レシピ、ランチャー、変換ツールはMaxTextに公開されている。