GoogleとMetaがTorchTPUをOSS公開 — GPU向けPyTorchコードが「置き換え2か所」でTPU上で動くように

11月3日、Sergio De Simoneが「Google and Meta Open-source TorchTPU to Bring Native PyTorch Support to Google TPUs」と題した記事を公開した。GoogleとMetaが共同開発したオープンソースプロジェクト「TorchTPU」によって、GPU向けに書かれたPyTorchコードがデバイス名と分散バックエンドの置き換え2か所だけでGoogle TPU上でネイティブに動作するようになった。既存のPyTorch資産をほぼそのまま活かせるという設計思想は、GPU以外のアクセラレータへの移行コストに悩む開発者にとって注目に値する。

PyTorch/XLAの限界を超える新アプローチ

GoogleはTPU上でPyTorchを動かす取り組みを2020年から開始しており、PyTorch/XLAという仕組みを使っていた。しかしこの方式には大きな問題があったとGoogleは述べている。

PyTorch/XLAは、コンパイル可能な領域を明示的に定義するためにソースコードへの大がかりな変更が必要で、コンパイル時のレイテンシが深刻で、PyTorchネイティブの命令実行（eagerモード）をサポートしておらず、PyTorchエコシステム内の分散コンピューティングの新しいプリミティブへの対応が遅れていた。

これに対しTorchTPUは根本的に異なる設計を採用している。eagerモードの実行を最優先とし、既存のPyTorchワークロードをほぼコード変更なしに動かすことを目指した。

移行コストの低さが最大の特徴

GPU向けに書かれたコードをTPU向けに移行する際に必要な変更は、2か所の識別子置き換えだ。cudaデバイスをtpuに、分散バックエンドのncclをtpu_distに置き換えるだけである。あるいは新しいtorch.accelerator.current_accelerator APIを使えば、利用可能なアクセラレータを自動的に識別でき、GPU/TPUを問わず同一コードで動作させることもできる。

import torch
import torch.distributed as dist

def main():
    device = torch.device("cuda") # -> device = torch.accelerator.current_accelerator
    dist.init_process_group(backend="nccl") # -> dist.init_process_group()
    rank = dist.get_rank()
    word_size = dist.get_world_size()
    ...
    model = model.to(device)
    ...

TorchTPUはeager実行モードを3種類サポートしている：strict、debug、DeferAndFuseだ。このうちDeferAndFuseはeagerとコンパイル実行のハイブリッドモードで、「Strict Eagerと比較して50%以上高速になることがあり、ピーク時のHBM（High Bandwidth Memory）使用量も削減できる」とGoogleは述べている。

約5,000モデルで動作検証

TorchTPUが既存のPyTorchモデルを無修正で動かせるかを確認するため、GoogleはHugging Faceの約5,000モデル（500以上のアーキテクチャ）を対象とした統合テストパイプラインを構築した。テストはTPUのv5e世代以降のハードウェアで、eagerモードとコンパイルモードの両方に対して実施している。

パフォーマンス検証では、300億パラメータのMuse Glimmerモデル（Googleが開発した大規模言語モデル）を標準的なPyTorchで実装して使用した。TPUv7チップ1枚での推論スループットは、シーケンス長5,120で25.7kトークン/秒、シーケンス長10,240で22.4kトークン/秒に達した。学習では4枚のTPUv7チップを使用し、それぞれ5.6k・4.6kトークン/秒のスループットを記録している。

また、Qwen3 MoE 235Bを用いたスケールアップ実験も行われており、TPUv7チップ1,024枚から8,192枚まで学習ワークロードをスケールさせた。推論については、大規模なMoEワークロードをvLLMで実行した際のコストパフォーマンスが競争力のある水準だったとしている。

Hacker Newsでの反応

Hacker Newsユーザーのtingkuanpeiはこのアナウンスに対し次のように述べた。

非常に興味深い取り組みだ。Googleはもっと前にこれをやるべきだった。PyTorchのエコシステムはあまりにも豊かだ。

また別のユーザーらは、GoogleがTPUをGoogle Cloud Platform（GCP）経由で提供するだけでなく、オンプレミス向けに企業へ直接販売していることも指摘している。顧客にはAnthropicが含まれており、Metaも将来的に加わる可能性があるという声も出ている。

今後のロードマップ

TorchTPUは現時点でまだ初期段階にあり、継続的に進化している。今後の開発計画には動的シェイプ対応、コンパイルパフォーマンスの改善、演算子カバレッジの拡充、強化学習への対応、量子化が含まれている。

詳細はGoogle and Meta Open-source TorchTPU to Bring Native PyTorch Support to Google TPUsを参照していただきたい。