この実験が重要な理由

大規模言語モデル(LLM)の学習は、もはや一部のビッグテックだけの領域ではありません。しかし「我々のモデルは順調に学習している」という主張は、多くの場合学習損失(training loss)カーブ一枚で証明されます。問題は、このカーブが嘘をつくことがあるという点です。

Google Cloud TPUチームは、Allen Institute for AI(Ai2)のOlmo 3 7BをMaxText(JAX/XLAベースのTPU学習フレームワーク)でゼロから再現しました。この実験が興味深い理由は三つあります。

  • 真のオープンモデル: 重み、データ、学習レシピ、ログまで完全公開されたフロンティア級モデルは極めて稀です。
  • フレームワーク間検証: PyTorch+GPUのレシピをJAX+TPUへ移した際、数学的に同一の結果が得られるかを確認できます。
  • 実務の落とし穴を公開: 学習損失は良好に見えるのに、実際の汎化性能は変わらないデータバグを発見したプロセスが克明に記録されています。

本記事は単なる論文要約ではなく、数週間にわたり実際にクラスタを回して得られた知見を実務視点で再構成したものです。根拠資料はGoogle Developers Blog 原文で確認できます。

核心メッセージ: 「再現した」と主張するには、学習損失ではなくホールドアウト(held-out)指標で証明する必要があります。

Large-scale language model training on TPU server cluster with distributed compute nodes Software Concept Art

再現の技術的核心: 何を合わせ、何を単純化したか

Olmo 3 7Bは32層、4096次元のdense transformerです。特徴はreordered normブロック、QK-norm、スライディングウィンドウ3:1 + グローバルアテンションの組み合わせです。MaxText設定(olmo3-7b-pt.yml)はこの構造を正確にミラーリングしています。

出発点: 重み変換からの検証

Ai2のstep-0 PyTorchチェックポイントをOrbaxへ変換後、順伝播の結果をHuggingFaceリファレンスと比較しました。結果はKL ≈ 1.5e-3、top-10トークンのうち9個一致 — 「同じモデル、異なるフレームワーク」水準のノイズフロアです。

# チェックポイント変換後の検証概念図(擬似コード)
# PyTorch → Orbax 変換後にKL divergenceを測定
kl_div = compute_kl_divergence(maxtext_logits, hf_reference_logits)
# 期待値: ~1.5e-3 (フレームワークノイズフロア)
assert kl_div < 2e-3, "変換失敗: モデルが異なる"

データパイプラインの再現

Olmo-coreと同一に実装しました。

  1. 全ドキュメントをトークン化しEOSで連結
  2. 8192トークン単位で非重複スライス
  3. 固定シードでグローバルシャッフル
  4. n-gram反復フィルタ(32回超の反復インスタンスをマスク)

MaxTextのdataset_type=olmo_grain(Grainベース)がこれを担います。

意図的な単純化2点

  • LRスケジュール: Ai2は当初5Tトークン計画でしたが実行中に5.93Tへ延長したため、コサイン曲線が二本繋がっています。MaxTextは全区間に単一コサインを適用しました。
  • データ: 公開されたOlmo-3 mixを使用(内部実行比で<0.5%のトークン欠落)。

これらの単純化にもかかわらず、ホールドアウト指標の全てでマッチが維持されました。これが「run-to-runノイズ範囲内での再現」と表現する理由です。

AI language model architecture diagram showing transformer layers and attention mechanisms

実務で本当に学ぶべき点: 学習損失の裏切りとハードウェア最適化

事例1: 学習損失が嘘をついた瞬間

約0.9MステップからMaxTextの学習損失がAi2より低くなり始め、1.25M以降は平均−0.06まで開きました。カーブだけ見れば「我々の勝ち」に見えます。しかし:

  • ホールドアウトC4損失は同率(Δ −0.004)
  • ダウンストリーム精度はむしろAi2が僅差で優位

原因はGrainデータローダーの二重シャーディングバグでした。ShardOptionsをサンプラーへ渡す一方、サンプラーが既に内部的にシャーディングしていたため、インデックスストリームが32倍速く前進しました。結果としてstage-1は一度のクリーンなepochではなくPoisson(≈1)リサンプリングとなり、コーパスの約37%は一度も見られず、26%は二回以上反復学習されました。反復されたインスタンスで損失が人為的に低下したのです。

教訓: 学習損失カーブのみをモニタリングすると、データバグを「性能向上」と誤認します。必ずステップ整列したチェックポイントでホールドアウト評価を並行してください。

事例2: ヘッド次元を一つ変えて12%の無料性能向上

Olmo-3 7Bの基本設定は32ヘッド × 128 head-dimです。num_heads × head_dim = 4096なので、16ヘッド × 256 head-dimに変えてもパラメータ数(7.298B)とFLOPs(1565 TFLOP/step)は完全に同一です。

IronwoodのMXUが256×256 systolic arrayのため、head-dim 128ではQK matmulでアレイの半分が遊びます。256に合わせるとアイドルサイクルが消え:

項目基本(128)最適化(256)
スループット508 TFLOP/s/device571 TFLOP/s/device
MFU44.2%49.6%
パラメータ7.298B7.298B (同一)
FLOPs1565 TFLOP/step1565 TFLOP/step (同一)

長期学習を開始する前に必ず確認すべきタダの最適化です。

事例3: オプティマイザdtypeの静かな罠

weight_dtype=bfloat16に設定すると、Adamのm/vモーメントがmu_dtype継承で静かにbf16へ降格されます。bf16の~3桁仮数部が初期ウォームアップの微細な更新を毎回切り捨て、1000ステップで損失+0.93が累積しました。weight_dtype=float32(デフォルト)を維持すると格差が30倍縮まります。「なぜ再現できないのか」の最大要因でした。

事例4: トポロジー非依存のレシピ

JAX/XLAスタックの最大の利点はレシピがチップ数から分離される点です。グローバルバッチ(512インスタンス、4.19Mトークン/step)は固定、チップ数は流動的。占有されたクラスタで余った容量を何でも掴み、数学は同一に保てます。実際にstage-2はIronwoodからv5pへ世代まで跨ぎましたが、XLAフラグはそのまま、device typeのみ変更して57.4% MFUを記録しました(v5pピークの57%)。

事例5: 正確な再開(resume)の三重決定論

47,684ステップ学習中、16時間でホスト障害によりstep 19,627で停止しましたが、19,500チェックポイントから再開し失った127ステップを再学習した結果損失・パープレキシティ共にΔ=0.000。これが可能な理由は三層の決定論が積み重なっているからです。

  1. XLAコンパイル済みTPUプログラム
  2. Grain反復子状態チェックポイント
  3. Orbax復元

特にstage-2でgrain_worker_count=4時にstateless resumeが壊れることを発見し、GrainのGrainCheckpointHandlerでデータ反復子位置まで原子的に共に復元するよう修正しました。

注意事項と限界

  • ホールドアウト評価なしに再現主張は禁止: 学習損失オーバーレイは必要条件に過ぎず、十分条件ではありません。
  • アーキテクチャ調整はハードウェア依存: head-dim 256最適化はIronwood MXU(256×256)に特化したものです。他のTPU世代やGPUでは再検証が必要です。
  • Stage 3とポストトレーニングは未完了: 本記事はstage 1(事前学習)とstage 2(mid-training anneal)までのカバーです。SFT/RLはレシピのみ作成済みの状態です。
  • データ順序ノイズは完全除去不可: 最終損失オフセット~+0.006 natはバグではなく、フレームワーク間のデータ順序差です。ドメイン別パープレキシティを見るとwikitext103は−0.0005で同率、wikipediaのみ+0.0135で最大 — 一ドメインだけ突出すればバグ、全ドメインが均等にオフセットすれば順序ノイズです。

Cloud TPU infrastructure for reproducible machine learning training pipelines Coding Session Visual

実務適用のためのまとめ

このケーススタディから持ち帰るべき核心は三つです。

  1. 「再現」の定義をホールドアウト指標で固めましょう。 学習損失カーブのみを見るとデータバグが性能向上に偽装されます。ステップ整列したチェックポイントでホールドアウトC4損失、多ドメインパープレキシティ、lm-evalスイートを最低3つのランドマークで比較してください。

  2. ハードウェア-ソフトウェア共同設計には無料の昼食が存在する領域です。 head-dimをMXUサイズに合わせるだけで、パラメータ・FLOPs変化なしに12%のスループットを得ます。長期学習をコミットする前に必ずアブレーションを回してください。

  3. 再開正確性は決定論三層の和です。 XLAプログラム、データ反復子状態、オプティマイザ状態を原子的に保存しないと、静かにリサンプリングされたデータで学習することになります。

日本の開発コミュニティ(Qiita/Zenn)的文脈で補足すると、国産クラウドでもさくらインターネットの高火力やGMOインターネットのGPUクラウドなど、大規模学習環境が整備されつつあります。こうした環境で長期学習を回すなら**「チップ時間(chip-hours)」を予算単位とする習慣**を身につけてください。本記事でstage-1は約77k Ironwood chip-hours(≈3,200 chip-days)を消費しましたが、チップ時間はスライスサイズに依存しない固定値なので、予算算定がはるかに正確になります。

次のステップ学習方向

  • MaxText公式ドキュメントでJAX/XLA学習スタックに入門する
  • Grainデータローダーのチェックポイントハンドラパターンを習得する(大規模学習再現性の核心)
  • lm-eval-harnessでホールドアウト評価パイプラインを自分で構築してみる
  • Olmo 3の**stage 3(長文コンテキスト適応)とポストトレーニング(SFT/RL)**レシピは未実行のため、自分で回すのが良い次の課題です

あわせて読みたい記事

本コンテンツは、信頼性の高い情報源をもとにAIツールを活用して作成され、編集者によるレビューを経て公開されています。専門家によるアドバイスの代替となるものではありません。