グーグル、自社製AI半導体でPyTorchを直接動かす基本ソフトを無償公開

米Googleは、自社開発のAI半導体「Tensor Processing Unit（TPU）」向けネイティブPyTorchバックエンド「TorchTPU」をオープンソースソフトウェアとして公開した。PyTorchのモデルや各種演算をTPUデバイス上で直接実行できるようにするもので、ライセンスにはApache 2.0が適用されている。

TorchTPUは、各種PyTorch演算向けのカスタムATenカーネル実装を備えるほか、XLA計算の実行を最適化するコンパイルキャッシュ機能を搭載している。さらに「torch.compile()」との統合によるグラフ実行モードや、データ並列およびテンソル並列を用いた学習を高速化するため、複数のコンピュータや計算用チップに処理を分担させて同時に計算を進める手法。">分散学習、デバッグやベンチマーク用のユーティリティ群を支援する。モデルの実装例としては、LlamaやQwen、ResNet、DLRM、minGPT、UNetなどが用意されている。

動作環境としてはPython 3.11以上をサポートしており、PythonのパッケージリポジトリであるPyPIではバージョン0.2.0が公開されている。プロジェクトのビルドや依存関係の管理にはBazel（Bazelisk）が用いられており、対話型ノートブック環境であるMarimoを用いたチュートリアルも提供されている。Google Compute Engine（GCE）のTPU v6eインスタンスなどのTPU仮想マシン環境から利用できる。