10月28日、I Programmerが「Google And Meta Release PyTorch Backend For TPUs」と題した記事を公開した。GoogleとMetaが共同でTPU向けのネイティブPyTorchバックエンドをオープンソースとして公開したことを伝えている。
PyTorchユーザーがTPUをそのまま使える時代へ
MLエンジニアがGPU向けに書いたPyTorchコードをTPUで動かすには、これまで相応の書き直しが必要だった。今回公開されたPyTorch Backend for TPUs（GitHubリポジトリ名：torch_tpu）はその障壁を大幅に下げる。
最大のポイントは移行コストの低さだ。既存のPyTorchコードを .to("cuda") から .to("tpu") に変えるだけで動作するよう設計されている。ただし、ピーク性能を引き出すにはモデルのシェイプやシャーディングのチューニングが別途必要になる点は留意したい。まず動かすだけなら追加の書き換えはほぼ不要、というのが正確な理解だ。
旧来のPyTorch/XLAが抱えていた問題
TPU（Tensor Processing Unit）はGoogleが開発したAI・MLワークロード向けの専用チップで、もともとTensorFlowのために設計されたものだ。PyTorchからTPUを使う仕組みとしては、2020年に**PyTorch/XLA**が登場していた。
ただし、PyTorch/XLAには致命的な欠点があった。
- コンパイル時間が遅い
- PyTorchのeager実行モード（命令型の即時実行）に非対応
PyTorchのエコシステムがeager executionを中心に発展してきたことで、PyTorch/XLAとのギャップは広がる一方だった。
GoogleとMetaの役割と設計方針
今回のプロジェクトはGoogleとMetaの共同開発として位置づけられている。GoogleはGoogle CloudのTPUインフラ・ランタイム・XLAコンパイラ基盤を提供し、MetaはPyTorchコアの開発元としてPyTorch Foundationを通じたエコシステム整備・eager execution設計の知見を持ち込んでいる。両者の強みを組み合わせることで、PyTorchユーザーが違和感なくTPUを扱える設計を実現した形だ。
その設計思想を具体的に示すと、主な変更点は以下のとおりだ。
- eager executionを最優先とした設計。JITコンパイルと実行時最適化を組み合わせ、op-by-opの実行とコンパイル実行のパフォーマンス差を縮める
- JAX依存なし。ランタイムにJAXの依存関係を持たない
torch.compile()は従来のLazy Tensorではなく、**TorchDynamo** ベースに変更。性能はJAXコンパイルと同等水準とされる- bounded dynamismの導入により、シェイプ変更によるXLAの再コンパイルを回避
- DDP、FSDP、DTensorといったPyTorchの分散学習の仕組みを広くサポート
Hugging Faceとの統合と実運用実績
このTPUバックエンドはHugging FaceのTransformers v5、Diffusers、Kernelsライブラリとビルトインで互換性を持ち、Hugging Faceで最もダウンロードされている上位5,000モデルがそのまま実行できるとされている。
すでに本番環境での利用実績もある。Google Cloudの大規模顧客の推論・学習ワークロードにおいて実運用されているという。
GPUとTPUの使い分けという現実問題
GPUとTPUの競争という文脈で見ると、このリリースはMLエンジニアにとって現実的な選択肢を増やすものだ。NVIDIAのGPU供給制約が続く中、Google Cloud上のTPUをPyTorchのコードベースをほぼそのままで活用できるようになる意味は小さくない。
繰り返しになるが、「コードをそのまま動かす」ことと「最高性能を出す」ことは別の話であり、本番運用でピーク性能を求める場合はTPU固有のチューニングが求められる。その点を踏まえた上で、選択肢の一つとして評価するのが現実的だろう。
PyTorch Backend for TPUsはGitHubでオープンソースとして公開されており、現在も活発に開発が続いている。
詳細はGoogle And Meta Release PyTorch Backend For TPUsを参照していただきたい。
