TorchTPUとは｜GoogleとMetaが開発したPyTorchネイティブTPUバックエンドの性能と仕様

PyTorch Blog アメリカ

概要

GoogleとMetaの共同開発による、Google TPU向けのネイティブPyTorchバックエンド「TorchTPU」が公開されました。`.to(“cuda”)`を`.to(“tpu”)`に変更するだけでモデルを実行できる高い互換性と、eager実行の最適化による業界競争力のある推論・学習性能を実現しています。これにより、PyTorchユーザーはJAXなどの他エコシステムへの依存なく、TPUの計算リソースを最大限に活用することが可能になります。

詳細

主要成果

GoogleとMetaのオープンソース共同開発により、Google TPU上でPyTorchをネイティブに動作させるバックエンド「TorchTPU」が導入されました。最大の成果は、既存のPyTorchワークロードを極めて少ないコード変更（例：.to(“cuda”)から.to(“tpu”)への変更）で移行でき、そのまま動作させる「アウトオブザボックス」な実行環境を実現したことです。また、eager実行の最適化、torch.compileの統合、ネイティブなハードウェア加速を通じて、業界競争力のある推論および学習性能を提供します。

技術・詳細

TorchTPUは、以下の設計原則に基づいて構築されています。

• PyTorchエコシステムへの準拠： 実行時にJAXへの依存を排除し、堅牢なeagerモードを実装しています。また、Lazy TensorではなくTorch Dynamoベースのtorch.compile()を採用し、JAXのコンパイルと同等の性能を目指しています。さらに、Process Group APIを通じてTPUコレクティブを低レイヤーで提供し、DDP、FSDP、DTensorなどのPyTorch分散コンピューティング手法を幅広くサポートします。
• ユーザー体験の最適化： 実行時やコンパイラのエラーに対して、ユーザーが対処可能なアクションを含むメッセージを提供します。また、torch.profiler、DTensor、Rayといったコミュニティ既製のツールと統合されています。
• 性能の追求： 大規模ワークロードにおいてGPU上のPyTorchに対抗できるソリューションを構築しています。eager実行向けにJITコンパイルアプローチを導入し、演算ごとの実行とコンパイル済み実行の性能差を埋めるランタイム最適化を実装しました。また、形状変更によるXLAの再コンパイルを回避する「有界ダイナミズム（bounded dynamism）」や、事前コンパイル可能な動的TPUカーネルを探索し、プログラミングスタックを改善しています。

現在の実装状況として、新しい「DeferAndFuse」eagerモードを含むeager/compiledモードの両方をサポートし、主要な深層学習アーキテクチャに必要なコアATen演算を含む約450のAten opsに対応しています。また、JAX/PallasおよびHelion DSLによるカスタムカーネルの作成も可能です。

背景・業界の文脈

Google TPUは、その設計から性能、エネルギー効率、コスト面で大きな利点があり、大規模モデルの学習と展開に成功してきました。しかし、従来はTensorFlowやJAXのエコシステムと密接に結びついていました。2020年に導入されたPyTorch/XLAはLazy Tensorに基づいていましたが、PyTorchの急速な進化に追いつかず、ソースコードへの侵入的な修正が必要であったことや、コンパイル時間の遅延、eager実行モードの欠如、分散コンピューティングプリミティブのサポート遅延といった課題を抱えていました。TorchTPUは、PyTorchをJAXの動作方式に無理に適合させるのではなく、原則に基づいたネイティブなTPUバックエンドとして2025年中盤から開発が開始されました。

今後の展望

TorchTPUは現在も活発に開発が進んでおり、すでに一部のGoogle Cloudの大手顧客による推論および学習のプロダクション環境で利用されています。互換性の検証では、Hugging FaceのTransformers v5、Diffusers、Kernelsライブラリとの適合性が確認されており、LLM（Qwen, DeepSeek, Gemma, Llama, Mistral, GPT, OPT, BLOOM, Falcon, StableML）、音声（Whisper, Wav2Vec2, HuBERT, MMS）、視覚（DinoV2, SigLIP, CLIP, Swin Transformer, ConvNeXt, ViT, ResNet, DETR）、テキストエンコーダー（BERT, ALBERT, Electra, RoBERTa, DeBERTa, T5, MT5, Flan-T5）を含む、ダウンロード数の多い上位5,000のモデルでアウトオブザボックスの実行が可能です。

毎週の技術動向レポートを無料でお届け

各分野の分析レポートを読む価値があるかどうか一目で判断できるインフォグラフィックをメールで受け取れます。

📢 メールマガジンに無料登録（週刊・技術動向レポート）

ご登録いただくと、Troy-Technical から週刊で技術動向レポート（メールマガジン）をお届けします。

• 取得したメールアドレス・選択分野は配信目的にのみ使用します。
• 第三者へ提供することはありません。
• 配信はいつでも解除できます（各メール下部のリンクから）。

詳しくはプライバシーポリシーをご覧ください。

同意して登録へ進む ▶

登録は1分・いつでも解除できます