GoogleはTorchTPUをオープンソースとして公開しました。PyTorch開発者が、同社のTensor Processing Unit（TPU）上でモデルをネイティブに実行できるようになります。あわせてGoogleは、推論サービングエンジンのvLLMとSGLang、学習フレームワークのTorchTitan、Ray、Hugging Faceのライブラリでも、TPUをサポートすると発表しました。
最も恩恵を受けるのは、モデルをPyTorchで記述し、デフォルトでGPU上で動かしているチームです。Googleによると、こうしたワークロードを、同社独自のAIチップであるTPUへ、最小限のコード変更で移行できるといいます。この主張が実運用でも成り立つなら、GPUのコストを検討しているチームは、ソフトウェアを作り直さなくてもGoogleのハードウェアを試せます。
開発者が変更する点
TorchTPUは標準的なPyTorchのワークフローに沿っています。開発者はまず、各演算が呼び出された時点で即座に実行されるeagerモードでデバッグします。その後、モデルを最適化済みのグラフに変換するtorch.compileに切り替えて、本番環境に備えます。分散学習のプリミティブも標準で組み込まれています。
Hugging Faceは、切り替えに必要な作業を引数1つにまで減らしました。Transformers v5では、device_map=”tpu”を指定するだけでモデルをTPUに送れます。GoogleはTransformersのテストスイートを毎晩TPU上で実行しています。生成モデル向けのDiffusersライブラリもネイティブに動作し、複数チップにまたがるテンソル並列にも対応しています。
Metaは、主要な本番ユースケースにおいて、TorchTPUを通じてLLM、広告、レコメンデーションのワークロードをすでにTPU上で稼働させていると明かしました。MetaでPyTorchエンジニアリング部門のディレクターを務めるJana van Greunen氏は、「単一のハードウェアアーキテクチャに縛られるわけにはいきません」と述べています。
Googleが示す性能数値
TorchTPU上でのTorchTitanによる学習では、Qwen 3 8BでモデルFLOPs利用率が最大52%に達しました。最大8,192基のIronwood TPUにわたって、ほぼ線形にスケールしています。ほぼ線形のスケーリングとは、スループットがチップ数にほぼ比例して増えることを意味します。モデルFLOPs利用率は、チップが理論上持つ演算能力のうち、どれだけを有用な学習処理に振り向けられているかを示す指標です。52%なら、およそ半分が学習に使われている計算になります。
これとは別に、推論リクエストをクラスタ全体に振り分けるオープンソースのルーターであるllm-dでも成果が出ています。GoogleはRed HatおよびIBMとともに、llm-d 1.0を公開しました。Zenteiq.aiによると、Google Kubernetes Engine上でllm-dを使ったところ、TPU v6eチップ1,000基あたりの1日の出力が14億7,000万トークンから59億7,000万トークンに増え、約4倍になりました。Zenteiqはこの向上を、各リクエストをすでにそのコンテキストを保持しているレプリカへ振り分けるルーティングのおかげだとしています。
まだプレビュー段階の機能
2つの機能は引き続きプレビュー段階です。1つは、広く使われている強化学習フレームワークVeRLに対してGoogleが提供したTPUバックエンドです。もう1つはTPUエージェントスイートで、AIコーディングツールと連携し、PyTorchモデルのオンボーディング、診断、推論チューニングを自動化します。Googleはさらに、旧TensorFlow LiteのLiteRTをPyTorchエコシステムに加えました。スマートフォンなどのデバイス上でモデルを動かすためのものです。
TorchTPUのコードはGitHubで無償公開されています。
翻訳元: https://www.helpnetsecurity.com/2026/10/21/torchtpu-pytorch-on-tpus/
