GoogleがTorchTPUをオープンソース化、TPUでネイティブPyTorchが利用可能に

Googleは2026年10月20日、PyTorch Conference 2026で、Tensor Processing Unit（TPU）向けのネイティブPyTorchスタックであるTorchTPUを、GitHubでオープンソースとして公開したと発表した。同社はGoogle Cloudブログでそう述べた。このソフトウェアはMetaのPyTorchチームと共同開発されており、既存のPyTorchワークロードを最小限のコード変更でTPU上で実行できるよう設計されている。

Googleによると、今回のリリースにより、PyTorchでネイティブに利用できる選択肢にGPUとCPUに加えてTPUが加わり、vLLM、SGLang、llm-d、TorchTitan、Ray、TorchRec、Hugging Face、LiteRTなどのツールを通じて、学習、サービング、デプロイに対応する。開発者はイーガーモードでモデルをデバッグした後、torch.compileを使って本番環境レベルの性能に移行できる。これにより、TPUの利用効率を高め、メモリーのオーバーヘッドを削減できる。

MetaのPyTorch Engineering担当ディレクターであるJana van Greunen氏は発表の中で、TorchTPUによって、Metaは主要な機械学習システムチームと柔軟にハードウェアを共同開発しながら、主要なLLM、広告、レコメンデーションシステムのワークロードにも対応できると述べた。同氏は「これにより、PyTorchのワークロードをGoogle TPUにデプロイし、エンジニアリングチームに直接フィードバックを提供するとともに、誰もが必要とするエコシステムの柔軟性を構築できます」と語った。

Hugging Faceとの連携とサービングの統合

GoogleはHugging Faceと提携し、Hugging Face Hub上のオープンウェイトモデルをTorchTPU上でネイティブに実行できるようにしました。Transformers v5では、TPUへのモデルのデプロイは、渡すことで行われます。 device_map="tpu"また、DiffusersライブラリはTPU上でネイティブに動作し、イーガーモード、torch.compile、複数のアクセラレーターにまたがるテンソル並列処理に対応する。Googleはさらに、TPUインフラストラクチャ上でHugging Face Transformersスイートを毎晩継続的にテストするための専用CI/CDパイプラインを構築したと述べた。Hugging Faceの製品担当バイスプレジデントであるJeff Boudier氏は、この協業によって、ハードウェア選択を重視する同プラットフォームの方針が実践されると語った。

サービングについては、TorchTPU上で動作するPyTorchネイティブのvLLMバックエンドが、Inferactとの協業により開発された。Googleによると、カスタムPallasカーネルを使ってTPUハードウェア上で連続バッチ処理、量子化、マルチチップ並列処理に対応する一方、上流のワーカーモデル、スケジューラー、APIサーバーは維持されるため、既存の推論スタックを変更する必要はない。InferactのCEO兼共同創業者であるSimon Mo氏は、この統合によりTPUでvLLMを一級のサポート対象として利用できるようになると述べた。

RadixArkと共同開発された、SGLang向けのPyTorchネイティブなバックエンドは、構造化ワークロードとエージェント型ワークロードを対象としている。発表によると、SGLangはRadixAttentionを使ってリクエスト間でKVキャッシュを再利用するため、プロンプトの先頭部分が繰り返される場合に再計算を省ける。また、TorchTPU版では、HiCacheのI/Oバックエンドも兼ねる分離型転送バックエンドを採用している。RadixArkのCTO兼共同創業者であるBanghua Zhu氏は、TorchTPUによって同社の開発が加速しており、さまざまなシリコン上でSGLangのサポートを展開していると語った。

GoogleはRed Hat、IBM、500人を超える貢献者と25の組織からなるコミュニティとともに、llm-d 1.0をCloud Native Computing FoundationのSandboxプロジェクトとして公開する。発表では、llm-dはPyTorchのサービングエコシステムとKubernetesベースの本番デプロイをつなぐものと説明されている。ルーターはプレフィックスキャッシュを考慮したルーティングを行い、リクエストごとに、そのコンテキストをすでに保持しているレプリカへ転送する。また、計算負荷の高いプリフィルと、メモリー負荷の高いデコードを分離する。バージョン1.0では、マルチテナントのフロー制御、ルーターの高可用性、安定版v1 API、エージェント向けのセッション認識型ルーティングが追加される。また、Google Kubernetes Engineでは、共通のllm-dエンドポイントをGPUまたはTPUのどちらでも実行できる。

Zenteiq創業者のSashikumaar Ganesan教授は発表の中で、llm-dを活用したGKE Inference Gatewayにより、同社の合成データ生成における推論スループットの合計が4倍を超えて向上し、1,000基のTPU v6eアクセラレーターあたりの1日あたりの生成トークン数が14億7,000万から59億7,000万に増加したと述べた。

学習とオーケストレーション

分散学習向けのPyTorchネイティブなリファレンス実装であるTorchTitanには、Googleが本番環境対応と説明するTorchTPU統合が加わった。この統合はtorch.compile、TorchTitanのSimpleFSDP、専用のPallasカーネル、グラフレベルの最適化を基盤としている。TorchTitanは、Llama 3 70Bおよび405B、Qwen 3 8Bおよび32B、MuseGlimmer 30Bなどの密モデルの事前学習に加え、Qwen 3 235Bの混合エキスパートモデルとFLUX.1 12Bの拡散トランスフォーマー向けに、TorchTPUのレシピを提供する。Googleによると、この統合は最大8,000基のIronwood TPU（TPU v7x）で検証されている。また、TorchTitanを通じたTorchTPUの学習では、Qwen 3 8Bで最大52%のモデル演算性能利用率を達成し、最大8,192基のIronwood TPUにわたってほぼ線形のスケーリングを実現した。

Googleは、強化学習フレームワークVeRLに、TPUバックエンドとエンドツーエンドのリファレンスレシピも上流統合した。分離型のトレーナー・サンプラー構成を基盤とし、数学的推論データセットで実証されたGRPOレシピは、開発者が既存のハードウェア非依存のモデルコードやアルゴリズムコードを維持したまま、VeRLのワークロードをTPU上で実行できるようにすることを目的としている。

Rayは、単一のワークフロー内でCPU、GPU、TPUをまたいでPyTorchのデータ処理、トレーニング、強化学習、サービングをオーケストレーションするためのTorchTPUネイティブ対応を追加する。AnyscaleのチーフアーキテクトであるEdward Oakes氏は、TPUをRayの第一級アクセラレーターとして利用できるよう、両社が連携して取り組んだと述べ、「TorchTPUは、業界標準のスタックにGoogle TPUを直接組み込む」と語った。

従来のアーキテクチャと提供状況

TorchTPUは2026年4月7日、PyTorchのPrivateUse1インターフェースを基盤とする「Eager First」設計について説明したGoogle Developers Blogの記事で一般公開された。この設計では、共有コンパイルキャッシュを基盤に、Debug、Strict、Fusedの各eagerモードを提供する。当時Googleは、Fused EagerはStrict Eagerと比べて性能が50％から100％超向上し、TorchTPU上のtorch.compileはTorch DynamoでFXグラフを取得して、XLAとStableHLOを通じて低レベル化すると説明した。また、カスタムカーネルはPallasとJAXで記述される。記事によると、このスタックはDistributed Data Parallel、FSDPv2、DTensorを標準でサポートし、ランク間で実行が異なる場合にも対応できるよう設計されている。さらに、TorchTPUの2026年のロードマップ項目として、公開GitHubリポジトリの立ち上げも挙げられていた。

torch_tpuリポジトリは現在、google-pytorch GitHub組織の下で公開されており、Apache 2.0ライセンスが適用されている。READMEには、カスタムATenカーネル、XLA計算用のコンパイルキャッシュ、データ並列処理とテンソル並列処理に対応した分散トレーニング、torch.compileとの統合、デバッグおよびベンチマーク用ユーティリティ、さらにLlama、Qwen、ResNet、DLRM、minGPT、UNetのサンプルが記載されている。このプロジェクトはPython 3.11以降をサポートし、torch-tpu 0.2.0パッケージは2026年10月19日にPyPIで公開された。Linux x86-64向けに、CPython 3.11から3.14までのホイールが用意されている。ユーザーガイドはtorchtpu.readthedocs.ioで公開されている。またREADMEによると、プロジェクトのインタラクティブノートブックはTPU VM上で実行する必要があり、最も手軽な選択肢はv6e GCEインスタンスである。

同じ発表でGoogleは、旧称TensorFlow LiteのLiteRTがPyTorchエコシステムに加わると述べた。LiteRT-Torchライブラリは、PyTorchモデルを単一の litert-torch.convert() 呼び出しで変換し、変換後のモデルはAndroid、iOS、macOS、Windows、Linux、ウェブ、IoTデバイスで動作する。PyTorchモデルのオンボーディング、診断、推論性能のチューニングを自動化することを目的としたTPUエージェントスイートは、GitHubで早期プレビューとして提供されている。

この発表では、TPU上でPyTorchを実行するには、Google CloudのTPUドキュメントを参照するよう開発者に案内している。