概要
TorchTPU は、Googleの TPU 上で PyTorch モデルを動作させるためのバックエンドを提供。 カスタムATenカーネルや分散トレーニング、 torch.compile() との統合をサポート。 インストールやビルド方法、VSCodeでの開発環境構築手順を詳細に解説。 ユーザーガイドとインタラクティブノートブックも利用可能。 Apache 2.0ライセンス でオープンソースとして公開。
TorchTPU: PyTorch向けTPUバックエンド
- PyTorch のモデルや演算を Google TPU 上で実行可能にするバックエンドの提供。
- ATenカーネル のカスタム実装により、様々なPyTorch演算をTPUに最適化。
- XLA計算 の最適化を支える コンパイルキャッシュ の搭載。
- 分散トレーニング （データ並列・テンソル並列）をサポート。
- Llama, Qwen, ResNet, DLRM, minGPT, UNet など幅広いモデルのサンプル実装。
- torch.compile() との統合により、グラフモードでの高速実行が可能。
- デバッグ・ベンチマーク用ユーティリティ の提供。
ユーザーガイド・インタラクティブノートブック
- 公式ユーザーガイド: https://torchtpu.readthedocs.io/ にて提供。
- ドキュメントフォルダ内に Marimoノートブック （Jupyter類似）を収録。
- marimo edit file.py コマンドでノートブックを編集・実行可能。
- TPU VM （v6e GCEインスタンス推奨）での実行が必要。
インストール方法
- PyTorchのリリース互換性 に準拠し、 Python 3.11以上 に対応。
- Python仮想環境 の作成例:
python -m venv .venvsource .venv/bin/activate
- pipによるインストール:
pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpu
- ソースからのビルド:
bazel build --config=wheel_common //ci/wheel:torch_tpu_wheel- 最適化モード （-c opt）がデフォルト、デバッグや開発時は
--//:allow_unoptimized_wheel=Trueを追加。 - Remote Build Execution (RBE) 未使用時は
--config=no_rbeを追加。 - ビルド後、
pip install ../bazel-bin/ci/wheel/*.whl --extra-index-url https://download.pytorch.org/whl/cpuでインストール。
開発環境・依存関係
- Bazel によるビルド・依存管理を採用。
- Bazelisk （Bazelのラッパー）のインストール例:
go install github.com/bazelbuild/bazelisk@latestwget -O ~/bin/bazel https://github.com/bazelbuild/bazelisk/releases/download/v1.27.0/bazelisk-linux-amd64
- インストール後は PATH に追加し、
bazel infoで動作確認。
VSCode + clangdによるC++開発
- VSCode に clangd拡張 を導入し、C++コードナビゲーションを強化。
- Microsoft C++拡張 は競合防止のためアンインストール推奨。
- コンパイルデータベース 生成:
tools/setup_clangd.pyを実行- 初回は30分以上かかる場合あり、再実行は高速。
nox -s refresh_compile_commands -- --no-buildでも実行可能。
貢献方法・ライセンス
- PyTorch on TPU プロジェクトへの貢献歓迎。
- CONTRIBUTING.md 参照、 CLA署名 ・コードレビュー必須。
- Apache 2.0ライセンス で公開、詳細は LICENSEファイル 参照。
Google Cloud以外にTPUってあるの？