Google、Cloud TPU向けネイティブPyTorchバックエンド「TorchTPU」をオープンソース化

Googleはオープンソース化しました TorchTPU、Cloud TPU向けのネイティブPyTorchバックエンドです。これは MetaのPyTorchチームと共同開発 されたもので、2026年10月20日のPyTorch Conference 2026で発表されました。既存のPyTorchコードは標準の torch.device("tpu")を通じてTPUをターゲットにします。TPUはGPUやCPUと並ぶネイティブなPyTorchデバイスとして扱えるようになりました。Googleチームは、TorchTPUが PyTorch/XLAを置き換える と述べ、すべての新規ワークロードでの利用を推奨しています。

TL;DR

• 規模: モデルではなくソフトウェアスタック。4つの公開リポジトリ：TorchTPU、vLLM-TorchTPU、SGLang-TorchTPU、TorchTitan。TorchTPUは PyPIでv0.2.0、Python 3.11以上、Apache 2.0ライセンス。
• 動作環境: Google Cloud TPU。 v5e以降の世代でテストされ、8,192個のIronwood（TPU7x）チップへスケール。サービングはFP8、MXFP4、NVFP4のチェックポイントに対応。
• パフォーマンス: デバッグにはeagerモード、本番レベルのスループットには torch.compile を使用。
• 最良: Ironwood上でTorchTitanを使ったQwen 3 8Bの学習で52%のMFU（自己申告）。
• 最悪: あるファインチューニングのケースでは、デフォルトのeager実行で1ステップあたり11.49秒、完全にチューニングした場合の2.78秒と比較。

TorchTPUとは？

TorchTPUは、PyTorchモデルをTPU上でネイティブに実行する、GoogleとMetaによるオープンソースのPyTorchバックエンドです。PyTorchの PrivateUse1インターフェースを通じてプラグインとして組み込まれます。テンソルは通常のPyTorchテンソルであり、サブクラスやラッパーは不要です。デバイスごとに1プロセスで実行され、テンソルビューとインプレース変更を保持します。

前身であるPyTorch/XLAは、レイジーテンソルのグラフを記録していました。開発者はコンパイル可能な領域をマークするコードを追加することがよくありました。PyTorch/XLAはまた 純粋なSPMDコードのみをサポートしていた。TorchTPUは分岐のあるMPMD実行にも対応しています。

TorchTPUはどのようにPyTorchコードを実行するのか?

TorchTPUには 3つのeagerモード:

• があります。Strict(デフォルト): 標準のPyTorchセマンティクスで各演算子を非同期にディスパッチします。
• Debug: テンソルの境界チェックを行いながら同期的に実行します。
• DeferAndFuse: eager演算子の連続実行を融合(fuse)してHBMトラフィックを削減します。 TorchTPU論文 によると、Strictより50%以上高速になることがあるとのことです。

ピーク速度を得るには、 torch.compile(backend="tpu") がTorchDynamoでグラフをキャプチャします。StableHLOを経由してOpenXLAへ lowered され、OpenXLAがTPUの機械語コードを生成します。コンパイル済み実行ファイルはプロセス、ホスト、ジョブ再起動をまたいでキャッシュされます。

カスタムの Pallasカーネル は標準のPyTorch関数として呼び出せます。プレビルドの Tokamax カーネルや、 Helion のタイルカーネルも動作します。

分散実行には tpu_dist プロセスグループを使用します。DDP、FSDP2、DTensorをサポートし、TPUインターコネクト向けのトポロジ認識メッシュも備えています。

TorchTPUはどれくらい速いのか?

以下の数値はすべてGoogleおよびパートナーによる自己申告です。

• 学習スケール: TorchTitanは Qwen 3 8Bで最大52%のMFUに達しました。Googleは8,192基のIronwood TPUまでほぼリニアなスケーリングを報告しています。
• チューニングの階段: TorchTitanのファインチューニング事例では、ステップ時間が11.49秒から2.78秒に短縮されました。 DeferAndFuse 単独で 1.15 倍でした。Full torch.compile は 3.80 倍でした.
• モデルカバレッジ: Google は約 5,000 の Hugging Face モデルを eager モードとコンパイルモードで継続的にテストしています。
• フリートサービング: Zenteiq は 1,000 個の TPU v6e チップあたりの 1 日あたりの出力を 1.47B から 5.97B トークンに増やしました。この向上は GKE 上の llm-d ルーティングによるものであり、TorchTPU 単独の効果ではありません。

TorchTPU にはどのような統合機能が同梱されていますか?

• Hugging Face: Transformers v5 は TPU に次のようにデプロイできます: device_map="tpu"。Diffusers は eager、 torch.compile 、テンソル並列をサポートしています。
• vLLM-TorchTPU: は Inferact と共に構築されました。上流のスケジューラと API サーバを維持しつつ、カスタム Pallas カーネルを使用しています。
• SGLang-TorchTPU: RadixArkで構築されています。RadixAttentionのプレフィックス再利用を維持します。
• サービング機能: FP8 KVキャッシュ、FP8、MXFP4およびNVFP4チェックポイント。推測的デコーディングにはMTPとEAGLE-3が含まれます。
• TorchTitan: Llama 3 70B/405B、Qwen 3 8B/32B、Qwen 3 235B MoEおよびFLUX.1 12B向けのレシピ。
• その他: llm-d 1.0、Ray、TorchRec、 VeRL (プレビュー)、および LiteRT-Torch がデバイス上へのエクスポート用に提供されています。

TorchTPUはPyTorch/XLA、TorchNeuron、JAXとどう比較されますか？

TorchTPUはどのようにインストールしますか?

TPU VMにパッケージをインストールし、次にを選択します。 tpu デバイス:

# pip install torch_tpu
import torch
device = torch.device("tpu")
x = torch.randn((4, 4), device=device)

The リポジトリのREADME は torch_tpu==0.2.0 でCPU用PyTorchホイールインデックスを固定（pin）しています。これがないと、pipはCUDAビルドの torchを取得してしまいます。完全なガイドは torchtpu.readthedocs.io.

主要なポイント

• TorchTPUはApache 2.0のオープンソースで、Metaと共同構築されました。
• eagerファーストのPrivateUse1ベースの実行により、PyTorch/XLAを置き換えます。
• torch.compile(backend="tpu") はStableHLOとXLAを通じてlowerされます。
• TorchTitanはQwen 3 8Bで52%のMFUを達成し、8,192個のIronwoodチップへスケーリングしました。
• vLLM、SGLang、Hugging Faceが標準のPyTorchを通じてTPU上で動作するようになりました。

ぜひ GitHubリポジトリ, 論文, ドキュメント および Google Cloudの発表もご覧ください。このプロジェクトの研究成果は研究者に帰属します。また、お気軽に Twitter をフォローし、 150k以上のML SubReddit に参加して、 ニュースレターも購読してください。お待ちください！Telegramを使っていますか？ 現在、Telegramでもご参加いただけるようになりました。

本記事 GoogleがCloud TPU向けネイティブPyTorchバックエンド「TorchTPU」をオープンソース化 初出: MarkTechPost.