Google、TorchTPUをTorchTitanに統合したTorchTitan TPUを公開 PyTorchモデルのTPU最適化手法も解説

TorchTitan TPUはMetaのTorchTitanのフォークで、https://github.com/google-pytorch/torchtitan/ で公開されている。GoogleとMetaが共同で、TPUネイティブの実行パスをTorchTitan内に開発した。モデル定義はCUDA固有のコードを含まない単一デバイス向けのPyTorch（nn.Module）で、並列化は別レイヤーで差し込む設計になっている。構成ツリーを書き換えることで、アテンションをSplash Attentionに、最終射影とクロスエントロピーを融合Pallasカーネルに、MoEのエキスパート行列積をグループ化Pallasカーネルに置き換えられ、いずれもフラグで切り替える。

性能面の中核はSimpleFSDPとGraph Trainerである。SimpleFSDPはパラメータのシャーディングをDTensorのパラメトリゼーションとして表現し、all-gatherとreduce-scatterをトレース済みグラフ内の通常のノードにする。FLUXではこれによりJAXのMaxDiffusionの99%の性能に達した。Graph Trainerはmake_fxでフォワード、損失、バックワードの学習ステップ全体を1つのXLAプログラムとして捕捉し、層をまたぐscanやグラフ分割、アクティベーションのホストメモリへのオフロードを可能にする。functional AdamWをグラフに含めたQwen3-8Bでは51-52% MFUに達した。

最適化の手順はプレイブックとして整理されている。ベースライン確立（eager、compile、Graph Trainer）、Pallasカーネルとタイルサイズの調整、バッチサイズとメモリの最適化、XLAとlibtpuフラグの調整とプロファイルの順で、多くはレシピ上の設定変更で済むとされる。FLUX.1の12BモデルをTPU v6e-8で学習した例では、1ステップが11.49秒から2.78秒へ4.1倍高速化した。Qwen3 235Bは64チップから8,192チップまでのweak scalingで、チップあたりスループットの低下が8%にとどまった。ただしこの規模のデータ並列は学習レシピではなくスケーリング実験だと明記されている。

HuggingFace Transformersのモデリングバックエンドを使うと、Hubのidからモデルを構築し、SimpleFSDPやper-block torch.compile、Splash Attentionなどを自動適用できる。対応する500以上のモデルを同じトレーナーで動かせるが、MoEチェックポイントはエキスパート並列が未対応である。新規モデルの移植向けに、コーディングエージェント用のonboardingスキルもリポジトリに含まれる。ロードマップには、ML Goodput MeasurementやMulti-Tier Checkpointingとの統合、float8レシピの拡張、TPU 8tと8iへの対応、エキスパート並列と長コンテキスト学習が挙がっている。