Google、PyTorchのままTPUを動かす「TorchTPU」を公開　MetaではSGLang基盤が本番運用

TorchTPUは、標準のPyTorch API、eager実行、torch.compile、分散実行をサポートするオープンソースで、GitHubのgoogle-pytorch/torch_tpuリポジトリにApache License 2.0で公開されています。Google Cloudの公式ブログによると、Metaとの共同開発で「既存のPyTorchワークロードを最小限のコード変更でTPU上で動かす」ことが狙いです。Googleはこれと合わせて、Hugging Face Transformers v5では対応モデルをdevice_map="tpu"の指定でTPUに配置できるとも説明しています。

推論では、vLLMとSGLangの2つのOSS推論エンジンにTorchTPU向けのネイティブバックエンドが用意されました。vLLM-TorchTPUはvLLMの開発者の創業したInferactとGoogleが、SGLang-TorchTPUはSGLangの開発者の創業したRadixArkとGoogleがそれぞれ開発しています。Googleの説明によると、いずれも既存のワーカーモデル、スケジューラー、APIサーバーを保ったまま外部プラグインとして組み込む構成で、既存の推論スタックを変更せずにTPUへ載せ替えられる点が特徴です。vLLMプロジェクトの公式ブログは、vLLMにすでにあるモデル定義をそのまま使い、TPUでも変更なしで動くモデルが大半であると説明し、推奨ハードウェアをTPU v7x（Ironwood）、v6e・v5e・v5pを実験的としています。

学習では、PyTorchの分散学習の参照実装TorchTitanに本番品質のTorchTPU統合が加わり、Googleの測定でQwen 3 8Bの特定の学習構成において最大52%のModel Flops Utilization（MFU、計算資源の利用効率の指標）を示し、最大8,192基のIronwood TPUでほぼ線形のスケーリングと安定性を確認したと報告されています。強化学習のフレームワークVeRLにはTPUバックエンドとGRPOのレシピが上流化され、分散処理のRayもTPUを第一級のアクセラレータとして扱います。Kubernetes上で推論を分散するllm-d 1.0はRed Hat、IBMなど500人超のコントリビューター、25組織とともにリリースされました。

Google TPUでPyTorchを動かす経路自体は、PyTorch/XLAなど以前から存在しました。今回の新しさは、PyTorchのネイティブな実行とコンパイルに加え、推論・学習・分散の主要スタックとの統合が一通りそろったことです。AWSも同時期にTrainium向けのTorchNeuronを公表しており、開発基盤を作り直さずにハードウェアを乗り替えられる構図の競争が進んでいます。ただし、これでCUDAの優位が消えるわけでも、全PyTorchモデルが無修正で同等の性能になるわけでもありません。vLLM-TorchTPUは対応ハードウェアやモデルに制限があり、量子化やカーネルの互換性は個別の確認が必要です。

Googleはこのほか、スマートフォンや組込み機器向けにPyTorchモデルを1回の変換呼び出しでLightweight端末向け形式にするLiteRT-Torchの公開と、TPU上のオンボーディングや性能調整を自動化するTPU agent suiteの早期プレビューも発表しました。MetaのTPU利用の範囲や規模、調達条件は公表されておらず、Tom's Hardwareなどが伝えた「Metaのデータセンター内で稼働」「TPU購入」の部分は一次資料からは確認できていません。今後の焦点は、新モデルの公開直後にTorchTPUで実用になるまでの時間と、GPU・TPU・Trainiumを同一条件で比較する第三者の測定です。