TorchTPU PyTorch TPUガイド：2026年のインストールと移行方法

PyTorchコードをcudaからtpuへ移す作業自体は、それほど大がかりではありません。ただし、TPUへのアクセス環境やパッケージの選択、コンパイル、テンソル形状によって、実際にうまく動くか、十分な性能が出るかは大きく変わります。

TorchTPUでできること、できないこと

TorchTPUはGoogle TPU向けのネイティブなPyTorchバックエンドです。ホステッドモデルAPIでも、TPU VMの代替でもありません。PyTorchの発表では、イーガー実行、torch.compile、分散処理プリミティブ、約450のATenオペレーションへの対応が説明されています。また、公開リポジトリには、ATenカーネル、コンパイルキャッシュ、分散実行のサンプル、torch.compileとの統合が掲載されています。

移行の第一歩は、見慣れたコードの変更です。

# Before
model.to("cuda")

# First TPU port
model.to("tpu")

ただし、これは移植性を確認するためのテストであって、性能を保証するものではありません。PyTorchの発表でも、TPUでピーク性能を引き出せるかどうかは、モデルの形状、シャーディング、ハードウェアに適したオペレーションに依存すると説明されています。

インストール前に準備しておくこと

TorchTPUを活用できるのは、Google TPUが用意された環境です。リポジトリでは、ノートブックをTPU VM上で実行することを推奨しており、最初の環境としてはv6e Google Compute Engineインスタンスが最も始めやすいとしています。CPUのみのノートPCでもパッケージのインストールはできますが、TPU上での実行確認はできません。

公開READMEではPython 3.11以降をサポートし、torch_tpu==0.2.0を指定しています。既存のCUDAビルドや無関係な依存関係を引き継がないよう、独立した環境で試すのがおすすめです。

CUDAを引き込まずに`torch_tpu`をインストールする

まず仮想環境を作成し、CPU版PyTorchの配布元と、バージョンを固定したTorchTPUパッケージを指定してインストールします。

python -m venv .venv
source .venv/bin/activate

pip install torch_tpu==0.2.0 \
  --extra-index-url https://download.pytorch.org/whl/cpu

READMEの説明どおり、追加のインデックスを省略するとCUDA版のPyTorchがインストールされる可能性があります。TPU VMでは適切でない依存関係なので注意してください。

最初の移行では、公開済みのwheelを使うのが無難です。ソースからビルドする場合は、BazelまたはBazeliskに加えて、リポジトリで案内されている--config=wheel_common、必要に応じて--config=no_rbeのフラグが必要になります。

最小構成のPyTorchプログラムでバックエンドを確認する

次のコードをTPU VM上で実行します。

import torch

x = torch.randn(2, 2, device="tpu")
y = torch.randn(2, 2, device="tpu")
z = x @ y

print(z)
print(z.device)

tpuデバイス上で結果が返れば、テンソルの確保と行列積がTPUバックエンドで実行されたことを確認できます。デバイスを利用できない場合は、モデルのコードを変更する前に、VMのランタイムとパッケージのインストール状態を確認しましょう。

TorchTPUでは、デバッグ向けのイーガー実行と、コンパイル実行向けのtorch.compileを利用できます。PyTorchの発表によると、プロジェクト内のテストでは、DeferAndFuseモードが厳密なイーガーモードより50%以上高速になる場合がありました。ただし、これはプロジェクトの報告であり、あらゆるモデルに当てはまるベンチマークではありません。

実際のPyTorchワークロードを移行する

まずは最小限の変更で単一デバイスの動作を確認し、その後でTPU向けの分散設定を追加していきます。

torch.device("tpu")でデバイスを指定するか、対応している環境ではtorch.accelerator.current_accelerator()を使います。
モデルと入力テンソルを同じデバイスへ移します。
コンパイルや分散実行を有効にする前に、短いバッチを1回実行します。
計測中はバッチサイズとテンソル形状を安定させます。形状が変わるとコンパイル処理が発生する場合があります。
大規模な移行を決める前に、代表的なモデル経路と分散パターンをテストします。

分散実行するCUDAコードでは、一般にNCCLを使います。一方、TorchTPUの公開サンプルでは、次のようにTPU向けのプロセスグループバックエンドを指定しています。

import torch
import torch.distributed as dist

device = torch.device("tpu")
dist.init_process_group(backend="tpu_dist")
model = model.to(device)

分散実行の具体的な起動コマンドは、TPU環境とワークロードによって異なります。デバイス文字列を変更しただけで、シャーディング、プロセス配置、入力パイプラインまで自動的に設定されるわけではありません。

初回実行と2回目以降の実行では、結果が異なることがあります。TorchTPUには、繰り返し使うグラフを再利用できるコンパイルキャッシュが用意されているため、ウォームアップ時間と定常状態のスループットは分けて記録してください。

TorchTPUが実用段階にある領域と、まだ調整が必要な領域

Googleは、TorchTPUについてvLLM、SGLang、TorchTitan、Hugging Faceのツール群との統合を案内しています。推論に関する発表では、KVキャッシュの移動やサービング機能を含む、vLLMとSGLang向けのネイティブなTPU経路を説明しています。また、GoogleのCloudに関する発表では、TorchTitan上のQwen 3 8Bで52% MFUを達成した結果や、本番利用に関する主張が紹介されています。

一方で、発表の見出しだけでは見えにくい実際の制約もあります。

項目	ソースから確認できること	自分で確認すべきこと
インストール	Python 3.11以降、`torch_tpu==0.2.0`、CPU版PyTorchのインデックス	使用するVMイメージとwheelを問題なく解決できるか
移植性	イーガーモード、`torch.compile`、`tpu`デバイス、分散処理への対応	未対応オペレーションと、モデル固有のフォールバック
性能	約450のATenオペレーション、およびTorchTitan上のQwen 3 8Bでプロジェクトが報告した52% MFU	使用モデルのウォームアップ、形状、バッチサイズ、コスト
開発ワークフロー	公開サンプル、ノートブック、デバッグ・ベンチマーク用ユーティリティ	TPUへのアクセス、Bazelの複雑さ、利用できる可観測性
現在の不足点	機能のカバレッジはまだ拡大中	ワークロードが、まだ文書化されていない機能に依存していないか

TorchTPUを採用するか判断する基準

すでにGoogle TPUを利用でき、PyTorchモデルを動かしていて、学習や推論で共通のフレームワーク経路を使いたいなら、TorchTPUは今から試す価値があります。特に有望なのは、1行の移植だけでGPU相当の効率が得られると期待するのではなく、形状の調整やプロファイリングに取り組めるチームです。

成熟した互換性マトリクス、予測しやすいコールドスタート、あるいは自分のスタックに不可欠なのにTorchTPUではまだ公開されていない機能が必要なら、重要なワークロードの移行は待ったほうがよいでしょう。まずは1つのモデル、1つのTPUスライス、固定した評価セットで段階的に進め、インストール時間、初回コンパイル、定常状態のスループット、出力の正しさを分けて計測してください。

TorchTPU FAQ

TorchTPUはPyTorch/XLAと同じものですか？

違います。TorchTPUは、イーガー実行とtorch.compileを中心に設計された、比較的新しいネイティブPyTorchバックエンドです。一方、PyTorch/XLAは従来から使われてきたTPU統合経路です。既存のXLAコードを移行する場合は、別途検討が必要になります。

TPUのないノートPCにTorchTPUをインストールできますか？

パッケージをインストールすることはできますが、互換性のあるGoogle TPUインフラがなければ、TPU上での実行は検証できません。公式リポジトリの実用的なセットアップは、TPU VMを前提としています。

最初のTorchTPUのステップだけ遅かったのはなぜですか？

初回実行では、グラフのコンパイルやキャッシュへの登録が行われることがあります。後続のイテレーションは分けて計測し、再コンパイルの原因を調べる間は形状を安定させてください。

`cuda`を`tpu`に変えれば、同じ性能が出ますか？

いいえ。デバイス指定の変更で確認できるのは、あくまでコードが動くかどうかです。性能を比較する前に、形状、シャーディング、オペレーションの対応状況、ウォームアップをプロファイリングしてください。

TorchTPUは本番利用できる段階ですか？

Googleは本番環境での推論・学習利用を説明していますが、公開されているエコシステムのコンポーネントや機能の成熟度はそれぞれ異なります。導入を決める前に、実際のモデル、サービングエンジン、分散パターンをテストしてください。

TPU VMに指定のwheelをインストールし、2×2の動作確認を行ったうえで、ウォームアップ済みのワークロードを現在のGPU経路と比較してみましょう。こうすれば、コードの移植性、コンパイラのウォームアップ、オペレーションの対応状況、ハードウェアのコストを切り分けて評価できます。