google-pytorch/torch_tpu

これは何？

主な機能

• 01TPU向けのカスタムATenカーネル実装。
• 02コンパイル済みのXLA計算を再利用するコンパイルキャッシュ。
• 03分散学習向けのデータ並列化とテンソル並列化。
• 04グラフモード実行のためのtorch.compile()との連携。
• 05Llama、Qwen、ResNet、UNetなどのモデルを使ったサンプルコード。
• 06デバッグとベンチマーク用のユーティリティ。

ユースケース

• 既存のPyTorch学習コードをGoogle Cloud TPU上で実行する。
• TPUデバイス間で大規模モデルを分散学習する。
• TPUハードウェア上でPyTorchモデルの推論を提供する。
• TPU向けの一般的なモデルの実装例を学ぶ。

クイックスタート

1. 1

   Python仮想環境を作成する

   python -m venv .venv
2. 2

   CPU版PyTorchのインデックスを指定して、PyPIから最新の公開リリースをインストールする

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

使う前に知っておきたいこと

• !Google TPUハードウェアを対象としているため、TPU VMなどのTPUを利用できる場合にのみ役立ちます。
• !Python 3.11以降が必要です。
• !ソースからビルドするにはBazelが必要です。また、完全なwheelをビルドするにはRemote Build Executionの認証情報、またはno_rbeフラグが必要です。
• !メタデータ上のライセンスはNOASSERTIONと記載されていますが、READMEにはApache 2.0と記載されています。

トピック

よくある質問

torch_tpu とは？

PyTorchのモデルや演算をGoogleのTensor Processing Unit（TPU）上で実行できるPyTorchバックエンド。 TorchTPUはGoogleのTPU向けPyTorchバックエンドで、既存のPyTorchコードをTPUデバイス上で実行できます。さまざまなPyTorch演算に対応するカスタムATenカーネルに加え、コンパイル済みのXLA計算を再利用して繰り返しの実行を高速化するコンパイルキャッシュを備えています。

torch_tpu は無料で使えますか？ライセンスは？

torch_tpu は Other ライセンスのオープンソースです。商用製品で使う前にライセンス条件を確認してください。

torch_tpu は今もメンテナンスされていますか？

活発にメンテナンス中。torch_tpu の最終コミットは 2026年10月22日 です。

torch_tpu は何に使えますか？

既存のPyTorch学習コードをGoogle Cloud TPU上で実行する。 TPUデバイス間で大規模モデルを分散学習する。 TPUハードウェア上でPyTorchモデルの推論を提供する。 TPU向けの一般的なモデルの実装例を学ぶ。

torch_tpu のインストール方法は？

Python仮想環境を作成する: python -m venv .venv CPU版PyTorchのインデックスを指定して、PyPIから最新の公開リリースをインストールする: pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpu

torch_tpu はどこで入手できますか？

ソースコードとインストール手順は GitHub にあります：https://github.com/google-pytorch/torch_tpu