Google、PyTorchをTPUで動かす「TorchTPU」をオープンソース化

Googleは、PyTorch（AIモデルの構築・実行に広く使われているツールキット）を同社のTensor Processing Unit（TPU）で動かすソフトウェア層「TorchTPU」をオープンソース化した。MetaのPyTorchチームと共同開発したもので、既存の処理を書き換えるのではなく、比較的小さなコード変更でTPUに移行できるよう設計されている。

Googleによると、TorchTPUは、呼び出しに応じて処理を実行するためデバッグに向くeager実行と、本番環境向けのコンパイル実行の両方に対応する。開発者はデバイス設定を「cuda」から「tpu」に切り替え、TPU用の分散処理バックエンドを利用できる。別の方法として、PyTorchのアクセラレーター検出APIも案内している。TorchTPUはGitHubで公開されているが、動的な入力形状への対応、コンパイルの高速化、演算子の対応範囲、強化学習、量子化などは今後の課題として残っている。

開発者にとって何が変わる？

今回の公開により、PyTorch開発者がすでに使っているソフトウェアとTPUを組み合わせやすくなる。Googleによると、TorchTPUは推論サービングツールのvLLMとSGLangのほか、TorchTitan、Ray、TorchRecなどの学習・オーケストレーション関連プロジェクトにも対応する。Hugging Faceによると、Transformers v5ではdevice_map="tpu"と設定すればTPUを選べる。同社のDiffusersライブラリーも、コンパイルモードや複数のアクセラレーターへの処理分散を含むTPU実行に対応している。

これは、PyTorchをTPUで動かすためにGoogleが以前提供していたPyTorch/XLAとは異なるアプローチだ。Googleは、PyTorch/XLAではコードに大きな変更が必要で、コンパイルに時間がかかり、PyTorchのeagerモードにも対応していなかったとしている。Googleは、InfoQによると、500以上のアーキテクチャにまたがる約5,000のHugging Faceモデルを評価するテストパイプラインを構築した。幅広くテストされていることは示されているが、すべてのモデルや処理が変更なしで確実に動くという意味ではない。

ほかのハードウェアと比べると？

Googleは、eager実行とコンパイル実行を組み合わせたDeferAndFuseモードについて、Strict Eagerモードより少なくとも50%高速で、アクセラレーターのピークメモリー使用量も抑えられるとしている。別のテストでは、標準のPyTorchを使って300億パラメーターのMuse Glimmerモデルを実行した。InfoQによると、シーケンス長が5,120の場合、TPUv7チップ1基で推論スループットは毎秒25,700トークン、10,240の場合は毎秒22,400トークンだった。TPUv7チップ4基での学習では、それぞれ毎秒5,600トークンと毎秒4,600トークンを達成した。これらはGoogleが公表した結果であり、独立した比較によるものではない。

TechCrunchはTorchTPUを、同じ週に発表されたAmazonのTorchNeuronと並べて紹介している。TorchNeuronも、AmazonのTrainiumチップでPyTorchの処理を実行するためのものだ。どちらも、NvidiaのGPU以外を使いやすくすることを目指している。TorchTPUの公開によりPyTorchユーザーの選択肢は増えるが、処理を移行すればコストが下がる、あるいは現在の環境より高い性能が得られることを示すものではない。