Google系の google-pytorch/torch_tpu は、PyTorch を TPU で動かすためのバックエンドを公開しているリポジトリだ。単に「コードがあります」という話ではなく、TPU 上で PyTorch のモデルや演算を回すための実装、ビルド手順、開発環境の整え方までひとまとめにしているのが特徴だ。しかも README の情報だけでも、実運用に近いところまでかなり踏み込んでいることが分かる。研究用の試作ではなく、実際に使う人をかなり意識した公開物だと思う。
このリポジトリは、Google の Tensor Processing Units、つまり TPU 向けに PyTorch のバックエンドを提供するプロジェクトとして説明されている。README では、PyTorch のモデルや演算を TPU デバイス上で実行できるようにすることが目的だと明記している。中核になるのは TPU 向けの ATen kernel 実装で、PyTorch の内部演算に相当する部分を TPU 側に持っていく構成だ。あわせて、XLA 計算の実行を最適化するための compilation cache も入っている。XLA は TPU 系でよく使われるコンパイル基盤なので、ここがあると毎回同じ計算をやり直さずに済む。
用途としては、分散学習と推論の両方を視野に入れている。README には Data Parallelism と Tensor Parallelism の両方に対応すると書かれており、Llama、Qwen、ResNet、DLRM、minGPT、UNet といった例も挙がっている。さらに torch.compile() との統合もあり、graph mode での実行を想定していることが分かる。デバッグやベンチマークのためのユーティリティも用意されていて、単なるサンプル集ではない。
使い方の案内もかなり具体的だ。ユーザーガイドは torchtpu.readthedocs.io にあり、そこにはインタラクティブ notebook の静的版があるという。notebook は Marimo 形式で、marimo edit file.py で開ける。TPU VM、あるいは v6e の GCE instance で動かすのが簡単だとしている。
インストール面では、Python 3.11 以降がサポート対象とされている。まず仮想環境を作り、pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpu で公開版を入れる方法が示されている。0.2.0 は PyPI 上の最新公開リリースとして記されている。一方でソースからビルドする場合は Bazel を使い、bazel build --config=wheel_common //ci/wheel:torch_tpu_wheel を案内している。wheel_common を付けないと、XLA/MLIR backend の扱いが崩れて import 時に duplicate static registrations の問題が起きる、とかなり踏み込んだ注意書きまである。RBE、つまり Remote Build Execution を使う設定も含まれていて、認証がない場合は --config=no_rbe を後ろに足す必要がある。順番まで指定されていて、実際にビルドを回す人向けの README になっている。
開発環境の整備としては、Bazelisk の導入方法や、VSCode と clangd の組み合わせも書かれている。tools/setup_clangd.py で compile_commands.json を生成し、C++ の補完や定義ジャンプを有効にする流れだ。貢献については CONTRIBUTING.md に従うこと、CLA の署名やコードレビューの手順があることも案内している。ライセンスは Apache 2.0。GitHub 上では 4,465 commits、295 stars、15 forks、44 issues、220 pull requests という表示になっている。
まず感じるのは、これは「TPU でも動きます」という宣伝ではなく、かなり具体的な実装資産の公開だということだ。README に出てくる情報の粒度が、利用者目線というより開発者目線に近い。ビルド設定の順番まで説明しているのは、単なるサンプルではなく、壊れやすいところを抱えた本番寄りのソフトウェアであることの裏返しでもある。ここまで丁寧に書かれていると、逆に言えば「動かすだけなら簡単ではない」ことも透けて見える。
ただ、その複雑さは悪いことばかりではないと思う。TPU は GPU と違う都合が多く、PyTorch のような汎用フレームワークをそのまま持ち込むと、どこかで変換や最適化が必要になる。torch.compile() 対応や compilation cache の存在は、その面倒を利用者から隠すための工夫だろう。もしこれがなければ、TPU に合わせたコードを書き直す負担がもっと大きくなるはずだ。つまりこのリポジトリは、TPU を「特殊な研究機材」ではなく、PyTorch の延長として使わせるための橋渡しに見える。
README が面白いのは、例に出しているモデルがかなり幅広いことだ。Llama や Qwen のような大規模言語モデル系から、ResNet のような古典的な画像分類モデル、DLRM のような推薦系、minGPT のような学習用途の小さめのモデル、UNet のような画像生成・セグメンテーション系まで含まれている。これは「特定のベンチマークだけ速い」ではなく、ある程度一般的なワークロードを視野に入れているサインだと思う。
一方で、こうした例の豊富さは、ユーザーにとっては安心材料である反面、導入の難しさも示している。TPU は用途が決まれば強いが、学習・推論・分散配置の考え方が絡むと、機械学習フレームワークとしての使い勝手が一気に難しくなる。だからこそ、README が具体的なモデル名を並べているのは重要だ。抽象的な「対応しています」より、どんな仕事に向くのかが伝わるからだ。少なくとも、単一用途のデモではなく、複数の実践シナリオをまとめて面倒を見る意志は感じる。
個人的にいちばん目を引いたのは、ビルドまわりの説明の長さだ。普通の OSS なら「pip install で入ります」で終わりがちだが、このリポジトリでは --config=wheel_common の必要性や、--config=no_rbe を後ろに置く理由まで書いている。さらに、これを外すと import 時に duplicate static registrations の問題が起こる、とかなり具体的に警告している。こういう書き方は、裏で何度も似た失敗が起きてきたからこそ出てくるものだろう。
この手の説明があるのは、むしろ成熟の証拠だと思う。初期のプロジェクトは「使う人が少ないから説明も少なくていい」になりがちだが、ここではその逆で、壊れ方まで含めて共有している。とはいえ、導入障壁が高いのも確かだ。Bazelisk、RBE、clangd、Marimo、TPU VM と、周辺要素がかなり多い。だから、誰でも気軽に触れるとは言いにくい。実際には、TPU を持っている組織や、ML infra に慣れたチーム向けの色合いが強いのではないかと思う。
このリポジトリは、PyTorch と TPU の距離を縮める実験場としても読める。PyTorch はもともと CUDA 前提の文脈が強いが、ここでは CPU 用 wheel を明示し、TPU 側のバックエンドを別に提供している。これは、既存の PyTorch コード資産をできるだけ崩さずに別の計算基盤へ持っていくための工夫だ。ユーザーは torch.compile() や既存の API 群を使いながら、裏側だけ TPU に寄せられる。その方向性はかなり筋がいいと思う。
同時に、こうしたバックエンドはフレームワーク本体の変化に強く引きずられる。README で Python 3.11+ の対応や PyTorch の Release Compatibility Matrix を参照しているのも、その不安定さを前提にしているからだろう。つまり、便利さの裏側には保守コストがある。そこを正面から受け止めている点は評価できる。逆に言えば、このリポジトリの価値は、TPU の性能そのものより、「PyTorch を TPU に持っていくための継続的な摩擦をどこまで減らせるか」にあるのではないかと思う。
