TorchTPU：GoogleがTPU向けネイティブPyTorchをオープンソース化

Googleは2026年10月20日、自社のTPUチップ上でPyTorchのコードをネイティブに実行するバックエンド、TorchTPUをオープンソース化しました。これが重要なのは、TPUを使うための古くて難しい方法だったPyTorch/XLAに代わるものだからです。PyTorch/XLAは11月15日にアーカイブされます。Googleによると、既存のコードはデバイス名を変える程度の作業で移行できるはずです。

TorchTPUとは

TPUは、Googleが設計し、Google Cloudを通じて貸し出しているAIチップです。PyTorchは、AIモデルの構築と学習に最も広く使われているライブラリです。これまでPyTorchは、PyTorch/XLAを通じてTPUに届いていました。PyTorch/XLAは、実行前にモデルをコンパイル済みのグラフに変換します。

TorchTPUは代わりに、TPUを通常のPyTorchデバイスにします。チームはPyTorchブログの発表で、「ワークロードのTPUへの移植は、model.to("cuda")をmodel.to("tpu")に変えるのと同じくらい簡単であるべきです」と書いています。GoogleはこれをMetaのPyTorchチームと共同で、2025年半ばから開発してきました。

プロジェクトのGitHubリポジトリには、PyPI上のリリース0.2.0、Apache 2.0ライセンス、Python 3.11以降が記載されています。Crypto Briefingによると、このリリースは4つのリポジトリにまたがり、PyTorch Conferenceで発表されました。ブログが挙げる4つはTorchTPU、vLLM-TorchTPU、SGLang-TorchTPU、TorchTitanで、コアのバックエンド、2つのサービングエンジン、学習用ライブラリをカバーしています。

コードの実行方法

TorchTPUは、約450のATen演算を実装しています。ATen演算は、あらゆるPyTorchモデルを構成する中核的な数学関数です。コードはデフォルトで、通常のPyTorchと同じく「eager」モードで1行ずつ実行されます。ブログでは、3つのeagerモードを説明しています。

モード	有効にする方法	動作
Strict	デフォルト	各演算を書かれたとおりに実行する
Debug	TORCH_TPU_LAUNCH_BLOCKING=1	デバッグしやすいよう、各演算の完了を待つ
DeferAndFuse	TORCH_TPU_DEFER_AND_FUSE=1	演算をまとめて処理する。「Strict Eagerより50%以上高速」

さらに速度を上げるには、torch.compile(backend="tpu")でモデルを1つのTPUプログラムにします。カスタムのチップ向けコードは、JAXのPallasかHelionで書けます。ブログには、vLLM、SGLang、TorchTitan、Ray、TorchRecとの統合が挙げられています。ブログはまた、Google Cloudの最大級の顧客の一部が、すでに本番環境でTorchTPUを使っているとも述べています。

モデルの動作状況

現時点での数字は、Google自身のテスト結果だけです。ブログは次の結果を報告しています。

テスト	結果
実行したHugging Faceモデル	4,917（500以上のモデルファミリーから）
最初のテスト段階に合格	98%（4,917のうち4,821）
出力がGPUと一致	97%（確認した657のうち638）
64個のTPU v7チップでのQwen3 235Bの学習	PyTorchで1,828トークン/秒、GoogleのJAXスタックでは1,937

Crypto Briefingは、最大8,192個のTPU v7xチップでQwen 3 8Bを学習した際のモデルFLOPs使用率が52%だったというGoogleの数字を引用しています。この指標は、チップの最大演算能力のうち、モデルが実際に使っている割合です。

PyTorch/XLAは終了へ

PyTorch/XLAリポジトリの冒頭には現在、「このプロジェクトは非推奨となり、TorchTPUに置き換えられました」という告知があります。さらに、リポジトリは「2026年11月15日にアーカイブされる」とあり、これにより読み取り専用になります。その助言は、「新しいPyTorchのTPUワークロードはTorchTPUで始め、既存のPyTorch/XLAワークロードは移行する」ことです。最後に掲載されている安定版のインストール手順では、torchとtorch_xlaが2.8.0に固定されています。

同じ日、AmazonのTrainiumチップ向けTorchNeuronバックエンドがパブリックベータになり、同じXLA方式をやめました。どちらのクラウド事業者も、今では自社のAIチップを普通のPyTorchデバイスとして扱っています。

開発者にとっての意味

TorchTPUによって、TPUはデバイス文字列の1つになります。これでGoogleのチップは、GPUと同じ価格比較の土俵に乗ります。

• PyTorch/XLAからの移行を今すぐ計画してください。 11月15日以降、リポジトリは修正を受け付けません。実行中のジョブにはtorch_xla 2.8.0を固定し、READMEが最も簡単な出発点として勧めるv6eマシンでTorchTPUを試してください。
• まず2%を確認してください。 Google自身のテストでも、まだおよそ50モデルに1つが失敗します。移行を計画する前に、自分が使うモデルそのもので短い学習または推論のジョブを実行してください。
• カスタムCUDAカーネルの作業時間を見込んでください。 手書きのGPUコードはそのまま移せません。PallasかHelionで書き直す必要があるため、その作業をコストに含めてください。
• バージョンを固定してください。 これはリリース0.2.0です。requirementsファイルで固定し、アップグレードのたびに結果を確認し直してください。
• 自分でベンチマークしてください。 ここにある速度の数字は、すべてGoogleによるものです。自分のワークロードで、GPUや他のクラウドのチップと比べたトークンあたりのコストを比較してください。