GoogleがTPU向けPyTorchバックエンドをOSS公開
原題: Google open-sources TorchTPU, a native PyTorch back end for TPUs
なぜ重要か
NVIDIAのGPU依存を下げたい企業にとって、TPUをPyTorchで使いやすくする選択肢が増えることは、AIインフラの調達多様化に直結する。
Googleは、TPU上でPyTorchモデルを直接実行するネイティブバックエンド「TorchTPU」をオープンソースとしてGitHubで公開した。カスタムATenカーネル、コンパイルキャッシュ、分散学習のサンプルコードを含み、既存のPyTorchワークフローをTPUへ移植する際の摩擦を大幅に下げることを目的としている。公開時点でスター数は398、フォーク数は20。
Googleは「TorchTPU」をGitHubリポジトリ「google-pytorch/torch_tpu」として一般公開した。本プロジェクトはPyTorchのバックエンドAPIに準拠した実装で、ユーザーが既存のPyTorchコードをほぼそのままTPU上で動かせるよう設計されている。
アーキテクチャ面では、PyTorchの演算子ディスパッチ層であるATenに対するカスタムカーネルを実装しており、TPU固有の命令セットへの変換を担う。また、繰り返し実行されるモデルのコンパイル結果をキャッシュする機構も組み込まれており、起動レイテンシの低減が期待できる。
リポジトリにはBazelベースのビルドシステム、Dockerファイル、CI設定、そして分散学習を想定したサンプルコードが同梱されており、研究者や企業エンジニアが本番環境へ適用するまでの参照実装として機能する。ドキュメントはReadTheDocsで別途公開されている。
コミット数は4,466件に達しており、公開前から相当期間の内部開発が行われていたことがわかる。プルリクエスト221件、Issue 48件という数字も、すでに一定のコントリビューター活動が存在することを示している。
GoogleはこれまでTPU向けのPyTorchサポートとして「torch_xla」を提供してきたが、TorchTPUはXLAコンパイラを経由せずにTPUのネイティブバックエンドとして機能する点が異なる。これにより、XLAの抽象化レイヤーに起因するデバッグの難しさや一部の演算子非対応問題を回避できる可能性がある。ライセンスはApache 2.0で商用利用も可能。