Olmo 3 7BのMaxTextによるTPU上での事前学習の再現：大規模学習のケーススタディ

概要

Google CloudとAi2のチームが、PyTorchベースのOlmo 3 7BモデルをJAXベースのMaxTextへ移植し、Google Cloud TPU上で学習プロセスと性能を完全再現することに成功しました。

詳細内容

本記事は、Ai2のオープンモデル「Olmo 3 7B」をGoogle Cloud TPU上で再現した詳細な技術レポートです。PyTorchからJAXへのモデル変換、チェックポイントの互換性検証、そして5.9兆トークンの大規模学習における損失曲線と下流タスク精度の検証結果を網羅しています。特に興味深いのは「データローダーのバグが学習損失を低下させ、一見性能向上に見えたが、検証データでは改善していなかった」という、汎化性能評価の重要性を示す教訓です。また、TPUのMXU（256x256）に合わせてヘッド次元を256に調整することで、精度を維持したままスループットを12.4%向上させたハードウェア・ソフトウェアの共同設計事例や、異なるTPU世代間でのシームレスな移行、44.5%〜57.4%という高いMFU（モデル・フロップス利用率）の達成についても詳しく解説されています。

元記事を読む他のサマリーを見る