MetaとGoogle、TPU向け基盤「TorchTPU」を公開

米MetaのPyTorchチームと米Googleは共同で、GoogleのカスタムAIアクセラレーター「Google TPU（Tensor Processing Unit）」上でPyTorchモデルをネイティブかつシームレスに実行可能にする新しいオープンソースバックエンド「TorchTPU」を公開した。従来のPyTorch/XLAが抱えていた大規模なコード改修やコンパイル待ち時間、eager実行の制約を抜本的に解消し、モデルコードを .to("cuda") から .to("tpu") に切り替えるだけで動作する環境を実現している。

従来のPyTorch/XLAの課題を克服する「PyTorch Citizenship」

Google TPUは従来、TensorFlowやJAXエコシステムを中心に最適化されてきた。2020年に提供が開始されたPyTorch/XLAはLazy Tensor方式を採用していたが、コンパイル可能領域を明示的に指定する侵襲的なコード改修が必要であり、コンパイル遅延の長さやPyTorch本来の命令的（eager実行（コードの記述行ごとに即座に演算を実行する対話的な実行モード））への対応不足、最新の分散プリミティブとの乖離が課題となっていた。

これに対しTorchTPUは、「PyTorch Citizenship（PyTorch第一級市民）」を設計原則に掲げ、以下の特徴を備える。

• JAXランタイム非依存​: 実行時にJAXライブラリを介さず、PyTorchのディスパッチャ（PrivateUse1）とC++バックエンドを直結
• ​柔軟なeager実行​: 逐次実行の「Strict」、同期デバッグ用の「Debug」、複数Aten演算を遅延統合して高速化する「DeferAndFuse」の3モードを提供
• ​標準的なコンパイル統合​: TorchDynamoをフロントエンドに採用し、torch.compile(model, backend="tpu") を通じたStableHLO単一グラフ生成に対応
• Process Group層での分散統合​: c10d::Backendを継承したtpu_distを提供し、DTensorやFSDPv2、SimpleFSDPなどのPyTorch分散フレームワークと完全連携

多階層コンパイルキャッシュとカスタムカーネル・分散通信

TorchTPUは、XLAコンパイラの最大の課題であった頻繁な再コンパイルを防ぐため、3層の「Compilation Cache」を備える。

1. Tier-1（プロセスローカル・インメモリ）​: 単一プロセス内でC++オブジェクトとしてキャッシュ
2. Tier-2（マシンローカル・共有メモリ）​: 同一ホスト上の複数プロセス間でバイナリシリアライズ共有
3. Tier-3（永続ストレージ）​: Google Cloud Storageなどの外部ストレージに保存し、ノード再起動や別ジョブでも再利用可能

また、ユーザー定義のカスタムカーネルに対応するため、JAX/Pallas連携（@torch.tpu.pallas.jax_op）や高スループットカーネル群「Tokamax」、統一DSL「Helion」の3経路を用意した。さらに、TPU上のSparseCoreエンジンに分散集約通信（All-GatherやReduce-Scatter）をオフロードすることで、TensorCore（MXU）の行列積計算と通信の並行重複実行を可能にしている。

Hugging Face 5,000モデルでの検証とPod規模スケーリング

コミュニティにおける実用性を検証するため、Hugging Faceでダウンロード上位のオープンソースモデルを用いた大規模テストが実施された。

評価指標・カテゴリ	検証規模・定義	測定結果
​総検証モデル数​	LLM、音声、画像、マルチモーダル等のオープンモデル	4,917モデル​
​対象アーキテクチャファミリー​	Transformer、Diffusersなど異なるモデル構造	500以上​ のファミリー
Tier 1（人気モデル）パス率​	Qwen、Gemma、Mistral、Llama、Whisper等でのクラッシュなし実行	**98%**​（4,821 / 4,917）
Tier 2（代表モデル）数値一致率​	GPU実行時との数値的一致（Numerical Parity）	**97%**​（638 / 657）
Bounded Dynamism（動的形状）​	動的シェイプ対応での数値的一致率	**99%**​（834 / 839）

さらに、TorchTitanを用いた大規模分散学習では、Qwen3 MoE 235Bモデルを最大8,192基のTPU v7チップ（16,384基のTensorCore）で学習させた際、1スライスから8スライスまでほぼ一定のMFU（ハードウェアの理論最大演算性能に対する実効性能の割合）（27.1%〜28.8%）を維持する線形なスケーリングを実証した。また、vLLM-TorchTPUを用いた推論ベンチマーク（SemiAnalysis InferenceX）では、8K/1Kシーケンスにおいて100万トークンあたり2.01セントという極めて高いコストパフォーマンスを記録している。

日本の開発者への示唆と未確認事項

世界のAI開発シーンにおいて、GPUクラスターの供給制約やクラウド利用料金の高騰は共通の課題となっている。これまでGoogle Cloud TPUの利用にはTensorFlow/JAXの知識やPyTorch/XLA特有の移植工数が大きな障壁となっていたが、TorchTPUによって既存のPyTorch資産がそのまま動作する環境が整った。日本のAIスタートアップや研究機関にとっても、計算リソースの調達先をGPU以外へ迅速に分散できる選択肢となる。

なお、今回の初期公開バージョンでは、動的シェイプ（Bounded Dynamism）がデフォルトで有効化されておらず、torch.compileの実行時には一時的にdynamic=Falseが強制される仕様となっている。シンボリック整数（SymInts）を完全にサポートした動的シェイプ対応は、今後のロードマップ（6〜12か月以内）における課題とされている。

引用元：PyTorch