10月20日、Google Cloudが「TorchTPU is now open source: Run native PyTorch on TPUs」と題した記事を公開した。
PyTorchといえばGPU(あるいはCPU)という前提が、変わった。PyTorch Conference 2026にてGoogleは、TPU向けネイティブPyTorchスタック「TorchTPU」をGitHub上でオープンソース公開した。これまでTPUでPyTorchを動かすには、PyTorchのオペレーションをXLA(Accelerated Linear Algebra)計算グラフに変換する複雑なレイヤーが必要で、デバッグの困難さやeager modeが使えないことがエンジニアの大きな障壁となっていた。今回の統合により、TPU(GoogleがAI/ML処理向けに独自開発したアクセラレータ)はGPU・CPUと並ぶPyTorchのファーストクラスデバイスとなり、その障壁が取り除かれる。
重要なのは「既存コードをほぼ変更しなくてよい」という点だ。単なる実験的サポートではなく、vLLM・Hugging Face・TorchTitan・Rayといった主要フレームワークとの実用レベルの統合が一気に揃った形での発表となっている。
TorchTPUとは何か
TorchTPUはMetaのPyTorchチームと共同開発されたライブラリで、既存のPyTorchコードをほぼ変更なしにTPU上で動作させることを目的としている。
特徴は以下のとおりだ:
- eager modeでのデバッグが可能。開発中はeager modeで動作確認し、本番環境では
torch.compileでTPU最適化を適用するという、通常のPyTorchワークフローがそのまま使える。 - 標準PyTorch APIをサポートし、分散学習のプリミティブもネイティブで動作する。
- vLLM・SGLang・TorchTitan・Ray・TorchRec・Hugging Face・LiteRTといった主要フレームワークとの統合済み。
MetaのPyTorchエンジニアリング ディレクターのJana van Greunen氏は次のようにコメントしている:
「TorchTPUは大きな前進だ。コアMLシステムチームと柔軟にハードウェアを共同開発しながら、LLM・広告・レコメンドシステムの主要な本番ユースケースでGoogleのTPUにPyTorchワークロードをデプロイできるようになった。」
Hugging Faceとの統合:device_map="tpu" 一行で動く
エンジニアが即座に恩恵を受けやすいのが、Hugging Faceとの統合だ。Transformers v5では、TPUへのモデルデプロイが次の一行で完結する:
device_map="tpu"
DiffusersライブラリもTPU上でネイティブ動作し、eager mode・torch.compile・テンソル並列処理による複数アクセラレータへのスケールアウトをサポートしている。継続的な品質担保のため、Hugging Face TransformersスイートのTPUインフラ上での夜次テスト(nightly CI/CD)パイプラインも構築済みだ。
推論エンジンとの統合
vLLM:高スループット推論
TPUネイティブ推論を手がけるスタートアップInferactとの共同開発によるvLLMのTorchTPUバックエンドは、continuous batching・量子化・マルチチップ並列処理をTPU上でサポートする。既存のワーカーモデル・スケジューラ・APIサーバはそのまま使え、推論スタックの変更は不要だ。
SGLang:エージェント型ワークロードの効率化
SGLangの主要開発元であるRadixArkとの共同開発によるSGLang-TorchTPUは、RadixAttention(KVキャッシュをリクエスト間でツリー構造で再利用する技術)を活用して、繰り返されるプロンプトプレフィックスの再計算をスキップできる。マルチステップ生成のレイテンシとスループットを最適化する設計だ。
llm-d 1.0:Kubernetesネイティブな分散推論
Red Hat・IBMおよび500名超のコントリビュータ・25組織と共同でllm-d 1.0をリリースした。PyTorch推論ワークロードをKubernetes上で分散・オーケストレーションするためのフレームワークで、CNCF(Cloud Native Computing Foundation)のSandboxプロジェクトとして位置づけられている。
プレフィックスキャッシュ対応ルーティングにより、各リクエストをすでにコンテキストを保持しているレプリカに転送できる。GKE上ではTPUとGPUを共通のllm-dエンドポイントで扱えるため、ワークロードに応じたハードウェア選択が可能になる。
Zenteiq.aiはllm-dを活用した結果、1,000台のTPU v6eあたりの1日あたり生成トークン数が14.7億から59.7億へ、4倍超の向上を報告している。
大規模トレーニングの実績
TorchTitan:8,000 TPUでの分散トレーニング
TorchTitanはPyTorchネイティブな分散トレーニングのリファレンス実装で、TorchTPUとの統合が完了している。torch.compile・SimpleFSDP(PyTorchのFully Sharded Data Parallelismの簡易実装)・Pallasカーネル(TPU向け低レベル演算カーネル)を組み合わせた最適化が利用可能だ。
最大8,000台のIronwood TPU(TPU v7x)での動作検証済みで、Qwen 3 8Bのトレーニングでは**MFU(Model Flops Utilization:理論ピーク性能に対する実効利用率)52%**を達成。これはTPU規模のクラスタとしては良好な数値であり、8,192 TPUまでほぼ線形なスケーリングを確認している。対応アーキテクチャはLlama 3(70B/405B)、Qwen 3(8B/32B/235B MoE)、FLUX.1(12B)などだ。
VeRL:強化学習(RL)ワークロード
VeRLは、LLMの後学習(post-training)における強化学習ワークフローを効率化するフレームワークで、DeepSeekをはじめとするフロンティアモデル開発で採用が増えている。今回、TPUバックエンドをアップストリームに取り込んだ。GRPO(Group Relative Policy Optimization)の参照レシピも提供されており、既存のハードウェア非依存なモデルコードを変更せずTPUへ移行できる。
Ray:異種ハードウェアのオーケストレーション
RayにネイティブTorchTPUサポートを追加した。これにより、データ処理・トレーニング・RL・推論を含むAIパイプライン全体を、CPU・GPU・TPUをまたいで単一ワークフローで管理できる。
エッジデバイスへの展開:LiteRTがPyTorchエコシステムに参加
旧TensorFlow Liteである**LiteRT**が、PyTorchエコシステムに加わった。LiteRT-Torchを使えば、litert-torch.convert()一行でPyTorchモデルをエッジデバイス向けに変換できる。変換後のモデルはAndroid・iOS・macOS・Windows・Linux・Web・IoTデバイス上で動作し、CPU(XNNPack)・GPU(ML Drift)・NPUへのアクセラレーションを単一APIで利用できる。クラウドからエッジまで、PyTorch一本で完結する環境が整いつつある。
詳細はTorchTPU is now open source: Run native PyTorch on TPUsを参照していただきたい。