PyTorchモデルをTPUへ移行する前に把握したい、TorchTPUの演算子・動的形状の制約とコスト・性能の評価軸です

何が起きたか：TorchTPUのオープンソース公開と開発体制の基本情報

GoogleとMetaのPyTorchチームは共同開発体制のもと、Google TPU（Tensor Processing Unit）向けのネイティブPyTorchバックエンド「TorchTPU」をオープンソースとして公開しました[1, 2]。公式発表は2026年10月20日に行われ、公式リポジトリおよびPyPIパッケージは google-pytorch/torch_tpu および torch_tpu の名称で管理されています[1, 2]。PyTorchコード内ではデバイス名として tpu を指定して呼び出します[1]。また、PyPIにおける初版パッケージ torch_tpu 0.2.0は公式発表前日の2026年10月19日にリリースされました[3]。本ソフトウェアのライセンスには Apache License 2.0 が適用されています[4]。

現行バージョンにおいてTorchTPUの公式FAQが積極的なサポート対象として明記しているTPU世代は、TPU v5e、TPU v5p、Trillium（TPU v6e）、Ironwood（TPU v7x）の4世代です[5]。次世代にあたるTPU v8以降のチップについては、検証および最適化が現在進行中と記載されています[6]。

ソースコードはオープンソースライセンスに基づいて公開されていますが、実行基盤はGoogle TPUのハードウェアおよび専用ランタイムである libtpu に結び付いており、Pallasなどの最適化機能もTPUを前提としています[7]。そのため、オープンソース化によってフレームワーク層のソフトウェア利用や改変の自由度は高まるものの、ハードウェアやクラウドインフラ基盤への依存そのものを解消するわけではないという技術的な境界が存在します[7]。

なぜ今か：PyTorch/XLAの構造的課題と推論・多様化ハードウェアへの要請

従来のTPU向け実行環境であった「PyTorch/XLA」は、いくつかの構造的な課題を抱えていました。PyTorch公式ブログの言及によると、従来のスタックはコンパイル遅延（compile-time latency）が大きく、PyTorch本来の命令的かつ逐次的なeager実行モードをサポートしていませんでした[8]。PyTorch/XLAのRFC（Request for Comments）でも説明されているように、LazyTensorと xm.mark_step() による明示的な中間表現（IR）グラフの追跡は、eager実行を基本とするPyTorchの標準的な開発体験と乖離し、デバッグの困難さや専用API群の習得コストを生んでいました[9]。さらに、従来のPyTorch/XLAは全並列ワーカーが同一処理を行う純粋なSPMD（Single Program, Multiple Data）コードのみを前提としており、並列実行の柔軟性に制約がありました[10]。

こうした課題の一方で、生成AIやマルチモーダル、エージェント型ワークロードの急速な普及に伴い、推論サービングの効率化が重要性を増しています[11]。開発現場からは、vLLMやSGLangなどの推論スタックをPyTorch上で直接カスタマイズし、デバッグを行いたいという強い要望が寄せられていました[11]。加えてMetaのPyTorch Engineering責任者は、単一のハードウェア構成に縛られず、既存のPyTorchワークロードをGoogle TPU上の本番システムへ展開する選択肢を確保する重要性を表明しています[12]。

TorchTPUは、JAXランタイムへの依存を排除した設計となっており、PyTorchネイティブの実行経路を提供しつつ、低レベルのバックエンドとしてXLAコンパイラスタックを活用します[13]。PyTorch/XLAの公式ドキュメントおよび公式FAQにおいて、同プロジェクトは非推奨（deprecated）と位置付けられ、TorchTPUへの置き換えが明記されました[14, 15]。現在では新規のTPUワークロードをTorchTPUで開始し、既存のPyTorch/XLAワークロードをTorchTPUへ移行することが公式に案内されています[15, 16, 17]。

補足：SPMDとMPMDの違いとは

SPMD（Single Program, Multiple Data）は、並列環境内のすべてのワーカーが同一の命令列を異なるデータに対して実行する計算モデルです。MPMD（Multiple Program, Multiple Data）は、ランクやワーカーごとに異なる処理フローや命令を実行できる計算モデルを指します。

何が変わるか：ディスパッチからTPU実行に至る処理パイプラインの仕組み

TorchTPUは、PyTorch本体の外部バックエンド（out-of-tree backend）としてディスパッチャの PrivateUse1 機構に登録され、dispatch key tpu を介して呼び出されます[18]。Pythonコード上で演算が呼ばれると、TorchTPUは内部でStableHLO builderを含む DeferredOp インスタンスを生成し、アクティブなストリーム（active stream）のキューに積んで後から実行します[19]。したがって、Pythonでの演算呼び出しとTPUハードウェア上での実際の計算開始には時間差が生じます[19]。

実行モードとして、TorchTPUは主に3つの経路を備えています[20, 21]。

1. Strict eager（既定モード）：ATen演算子を1つずつ個別に非同期ディスパッチしてTPU上で実行します[20]。
2. DeferAndFuse（Fused Eager）：連続する演算列を実行時にオンザフライで融合（fusion）し、計算密度の高い大きなチャンクにまとめてTPUへ送出します[20]。
3. torch.compile：TorchDynamoが計算グラフを取得し、TorchTPUがそのグラフをStableHLO形式へ変換した上で、OpenXLAを用いてTPU向けのマシンコードへとコンパイルします[21]。

開発者は既存のPyTorchモデルをまず既定のeagerモードで動かして挙動の確認やデバッグを行い、その後に torch.compile を適用してTPUの利用率向上やメモリ削減といった本番向け最適化を進める開発ワークフローを採用できます[22]。

このパイプラインは、従来のPyTorch/XLAやGPU向け実行スタックと構造が異なります。従来のPyTorch/XLAが演算をIRグラフに蓄積し、xm.mark_step() などの同期タイミングでまとめて実行していたのに対し、TorchTPUは個別演算の非同期ディスパッチを標準としています[23]。また、一般的なGPU向けの torch.compile はTorchInductorとTritonを用いてGPUコードを生成しますが、TorchTPUはOpenXLAとStableHLOを経由してTPU専用コードを生成します[24]。TorchTPUはPyTorch/XLAの置き換えを担いつつ、コンパイル基盤としては下層のOpenXLAとStableHLOを継続して活用しています[25]。

補足：StableHLOとは

StableHLOは、機械学習コンパイラ向けの高水準演算中間表現（HLO）です。異なるフレームワークとハードウェアコンパイラ間の共通インタフェースとして機能します。

何が変わるか：演算子・動的形状・数値精度の対応範囲と実務上の制約

初期公開版（バージョン0.2.0）の公式ドキュメントによると、TorchTPUはforwardパスとbackwardパスの両方において450のcore ATen演算子をサポートしています[26]。ただし、すべてのPyTorch演算が無条件に動作するわけではありません。個別演算の制約として、現時点でFlexAttentionのネイティブ対応は提供されておらず、view演算はTPU上で計算オーバーヘッドが大きくなる可能性があります[27]。

公開されている機能対応表（feature set）では、Qwen、Llama、Muse Glimmer、GemmaなどのDenseおよびMoE Transformerモデルが検証済み（Validated）に分類されています[28]。一方で、GLM、Kimi、DeepSeek、GPT OSS、Mamba、LSTM/GRU、DLRMなどのモデルやSymIntsは実験的（Experimental）な扱いです[28]。さらに、FSDPv2やMPMDには低性能（low perf）の注記が付与されています[28, 29]。

動的形状（Dynamic Shapes）を扱う処理にも明確な制約が存在します。TorchTPUは動的次元のコンパイル時上限を固定する「bounded dynamism（境界付き動的形状）」を採用しています[29, 30]。1つのコンパイルトレース当たり最大8バケットに制限されており、固定した上限を超える新たな形状が入力されると再コンパイルが発生します[30, 31]。また、実行時のデータ値に依存して出力形状が決まる演算（例：torch.nonzero）やPallasカスタムカーネルとの併用は未対応であり、dynamic=True オプションもサポートされておらず、指定すると dynamic=None の使用を促す例外が発生します[30, 31]。

数値精度に関して、AMP（自動混合精度）APIではbfloat16、float16、float32、float64、float8系、整数、boolean、complexなどのデータ型（dtype）が列挙されています[32]。ただしこの列挙は、個々の演算子・dtype・形状のあらゆる組み合わせでの動作を網羅的に保証するものではありません[32]。

MXU（Matrix Multiply Unit）におけるfloat32の行列積（matmul）計算には、精度と計算コストのトレードオフが存在します。既定の「medium」設定では、入力float32をbfloat16へと丸めた上でfloat32として累積計算を行います[33]。精度を高める設定として「high」および「highest」が用意されていますが、high設定は3-pass計算によりMXU計算コストが概ね3倍、highest設定は6-pass計算により概ね6倍に増大します[33, 34]。この倍率はMXU上の純粋な演算コストを指すものであり、モデル全体のステップ実行時間がそのまま3倍や6倍になることを意味するわけではありません[34]。

今回確認した公式資料の範囲では、演算子・dtype・形状・実行モード・TPU世代を網羅した詳細な互換性マトリクスや、PyTorch/XLAの全APIからの移行対応表は確認できませんでした[35]。450のcore ATen演算子という総数や代表的な動作確認モデルの存在のみを根拠として、個別のモデル構成が完全に互換動作すると一律に判断することはできず、各ワークロードに応じた事前検証が求められます[35]。

何が変わるか：公表性能・ベンチマークの測定条件と独立評価の実態

公式発表においてTorchTPUの大規模な動作検証結果が公表されています。4,917種類のモデルを対象とした検証では、4,821モデル（98%）がクラッシュすることなく実行されたと報告されています[36]。この98%という数値はモデルがエラーで停止しなかった実行可否を示す指標であり、数値精度の一致率を表すものではありません[36]。別枠で実施されたGPU基準の数値検証では、代表的なアーキテクチャ657モデル中638モデル（97%）が合格し、SymIntsを用いたbounded dynamism検証では因果LLM 839モデル中834モデル（99%）で数値の一致が確認されています[36]。数値検証の合格基準としては、因果LLMではTPUと基準側（GPU）のperplexity相対差が0.10以内、画像・音声・エンコーダーの高次元出力ではRelative MSE（相対平均二乗誤差）が0.10以下、または絶対MSEが0.05未満と定義されています[37]。

学習性能に関して、公式論文および技術ブログで具体的なモデル測定値が報告されています。純粋なPyTorch実装によるMuse Glimmer 30Bでは、TPUv7 1チップ（FSDP world size 2）でのprefillスループットが系列長5,120で約25.7k tokens/s、系列長10,240で約22.4k tokens/sを記録し、TPUv7 4チップのFSDP学習（activation checkpointing使用）では同じ系列長でそれぞれ約5.6k tokens/s/TPU、約4.6k tokens/s/TPUのスループットが報告されています[38]。また、GPU向けに記述されていたMuse Glimmer 30BをTorchTPUへ移植し、FlexAttentionを scaled_dot_product_attention に置き換えて torch.compile とSimpleFSDPを適用した例では、TPUv7 4チップで系列長5,120および10,240の学習スループットが約5.3k tokens/s/TPU（MFU 36%）および約4.2k tokens/s/TPU（MFU 29%）に達しています[39]。さらに別のTorchTitan事例では、TPUv7x 64チップを用いて系列長16,384、1 sequence/device、Graph TrainerとSimpleFSDP、FP32 gradient reductionを適用した構成において、5,848 tokens/s/chipのスループットと約44%のMFUが記録されています[40]。

超大規模モデルであるQwen3 MoE 235Bの初期スモークテストでは、64 TPUv7構成においてTorchTitan／TorchTPUが1,828 TPS、JAX／MaxTextが1,937 TPSを記録したと報告されています[41]。ただしこの記述では、TPSの分母（チップ単位かクラスタ全体か）、バッチサイズ、精度型、コンパイルやウォームアップの測定区間といった詳細な前提条件は特定されていません[41]。Google Developers Blogに掲載された詳細な弱スケーリング検証では、系列長4,096、16系列/device、BF16パラメータ、FP32勾配reduction、1,024チップslice内FSDPとslice間HSDPの構成を用い、TPUv7xを64チップから8,192チップへ拡大した際、チップ当たりのスループットは4,568 tokens/s/chipから4,208 tokens/s/chipへと約8%の低下に抑えられたと報告されています[42]。なお、8,192チップ構成におけるQwen3 235BのMFU値について、公式論文内の初期報告表では27.1%と記載されている一方、技術ブログ記事では30.51% MFUと記載されており、同一条件での更新値か算定基準の相違かが明示されていないため、単一の確定値として扱うことは避ける必要があります[43]。また、Qwen3-8BにおいてGraph Trainerとfunctional AdamWを用いて51–52% MFUを達成したという報告もありますが、TPU世代やチップ数、バッチサイズなどの測定環境は明記されていません[44]。

画像生成モデルのFLUX.1-dev（TPUv6e 8チップ、1024×1024解像度、BF16、FSDP、バッチサイズ3、full activation checkpointing）を対象とした検証では、最適化手法の段階的適用による実行速度の短縮推移が示されています[45]。開始時のeager＋FSDP2では11.49秒/stepを要していましたが、最終的な局所コード改修後には2.78秒/stepへと短縮され、初期状態と比較して4.1倍の高速化、同一ハードウェア上のJAX MaxDiffusion実装と比較して約99%のスループットを達成しました[45, 46]。

最適化フェーズ	ステップ実行時間（秒/step）	適用内容の概要	出典
初期構成	11.49	eager 実行 ＋ FSDP2（基準測定）	[45]
DeferAndFuse 適用	9.95	連続する演算列のオンザフライ融合	[45]
SimpleFSDP 適用	6.52	TPU 向けに最適化されたシャーディング管理	[45]
Pallas SplashAttention 適用	5.55	TPU 向けカスタムアテンションカーネル導入	[45]
per-block torch.compile 適用	3.02	ブロック単位のグラフコンパイル	[45]
モデルコード局所改修後（最終）	2.78	TPU 向けコード調整（初期比 4.1 倍速、JAX 比約 99%）	[45, 46]

推論デコード処理において、Hugging FaceはTPUv6e 8チップ環境、バッチサイズ64でHubカーネルとdefault実装を比較し、Qwen3-30B-A3Bで21.5倍、Gemma 4 26B-A4Bで12.8倍の速度向上を確認したと報告しています[47]。

外部機関による評価として、SemiAnalysisがInferenceXベンチマークを通じた推論評価を公表しています[48]。Qwen3.5 397B（FP8、入力8,000トークン／出力1,000トークン）の大規模サービングにおいて、TPUv7 IronwoodとNVIDIA B200およびB300を比較しています[48]。ユーザ当たり100 tokens/sの生成速度を満たす条件での外部TCOモデル試算では、100万total tokens当たりの費用がIronwoodで約0.181ドル、B200で約0.222ドル、B300で約0.276ドルと試算され、ユーザ当たり20 tokens/sの条件ではチップ当たりスループットがIronwoodで9,364 tokens/s、B200で8,903 tokens/s、B300で8,925 tokens/sと報告されています[48]。

一方で、同時リクエスト数（concurrency）256の環境における初回トークン生成遅延（mean TTFT）を比較すると、Ironwoodが5.41秒であるのに対し、B200は3.75秒、B300は2.40秒と報告されています[49]。トークン単価やスループット指標における効率性と、ユーザ応答の初動レイテンシ特性は異なる挙動を示しています[49]。

ハードウェア構成	推定総保有コスト（USD / 百万 tokens） ※100 tokens/s/user 条件	チップスループット（tokens/s/chip） ※20 tokens/s/user 条件	平均 TTFT（秒） ※concurrency 256 条件	出典
TPUv7 Ironwood	0.181 ドル	9,364 tokens/s	5.41 秒	[48, 49]
NVIDIA B200	0.222 ドル	8,903 tokens/s	3.75 秒	[48, 49]
NVIDIA B300	0.276 ドル	8,925 tokens/s	2.40 秒	[48, 49]

現在公表されている性能データは、JAX実装（MaxTextやMaxDiffusion）、GPU環境、あるいはNVIDIA GPUとのサービング比較が主軸となっており、同一ワークロードにおけるPyTorch/XLAとTorchTPUの直接的な性能比較値は公開資料で確認できません[50]。また、SemiAnalysisの評価は第三者視点による客観的分析としてPyTorch側からも紹介されていますが、著者自身がGoogle開発チームへの謝意を記載している通り、これはQwen3.5推論に関する外部評価であり、学習ワークロードの独立検証や再現実験とは性質が異なります[51]。

誰に影響するか：モデル学習・推論サービングの移行手順と分散構成

既存のGPU向けPyTorchコードからTorchTPUへの移行は、デバイス指定を .to("cuda") から .to("tpu") へ書き換える手順から着手できます[52, 53]。ただしPyTorch公式ドキュメントが注意を促している通り、デバイス指定の置換のみでTPUのピーク性能が得られるわけではありません[52, 53]。GPU向けに偏ったテンソル形状（shape）、シャーディング設定（sharding）、演算実装（operations）が存在する場合、TPUアーキテクチャに適した形への調整が必要となります[53, 54, 55]。

分散並行処理において、TorchTPUは専用のバックエンドである tpu_dist を提供しています[55, 56]。1つのTPUデバイスに対して1つのプロセスを割り当てるモデル（one-process-per-device）を採用しており、標準の torch.distributed APIをeagerモードおよび torch.compile の双方からそのまま呼び出すことが可能です[56]。

大規模な事前学習やファインチューニングでは、TorchTitanと torch.compile を組み合わせた実行経路が案内されています[57]。GoogleとMetaが共同開発したTPU向けシャーディングおよびコンパイルパス、そしてSimpleFSDPを活用してモデル層の分割と通信を最適化します[57]。公式機能表において、Dense/MoE Transformer（Qwen、Llama、Muse Glimmer、Gemma）はValidatedに分類される一方、その他の多様なモデルはExperimentalに位置付けられています[58]。また、ランクごとに処理が異なるMPMD実行やコレクティブ通信は機能としてサポートされているものの、現時点ではピーク性能に達しないため、学習性能を重視する場合は torch.compile とSimpleFSDPの利用が推奨されています[59]。

推論サービング領域では、vLLMやSGLangの標準的な構成を維持しつつ、推論時のモデル定義、投機的デコーディング、量子化方式などをPyTorch上で直接カスタマイズおよびデバッグできる環境が整備されています[60]。実運用の実績として、MetaのPyTorch-TPU責任者は SGLang-TorchTPU がMetaの本番環境で稼働していることを明かしています[61]。ただし、公開されている資料からはMetaにおけるTPUの配備規模、総処理量に占める割合、契約容量、GPUと比較した実測TCO削減額などは確認できません[62]。そのため、個別の実稼働事例から採用規模や経済的効果全体を一般化して推し量ることはできません[62]。

マルチホストの大規模クラスタ運用において、Google CloudはGKE上でRayおよびKubeRayを用いた構成例を示しています[63]。WebhookによってPodホスト名、ヘッドレスサービス（headless Service）、およびTPU環境変数がマルチホスト協調用に自動構成されます[63]。さらにRay Serveを用いた推論サービングでは、トポロジーを考慮したスライス配置（topology-aware slice placement）により、モデルレプリカごとに分断されていない無傷のTPUスライスを確保し、リクエスト同時実行数に応じてスライス単位でレプリカを増減する運用手法が採られます[64]。公式ドキュメントはPyTorch/XLAを非推奨とし、新規および既存のワークロードをTorchTPUへ移行することを推奨しています[65]。

誰に影響するか：環境構築・クラウド提供条件とCPU向けoneDNN統合との判断境界

TorchTPUを実行可能な環境として、Google Compute Engine上でプロビジョンされたCloud TPU VM、Google Kubernetes Engine（GKE）、およびGoogle ColabのTPUランタイムが提示されています[66]。

ソフトウェア要件に関して、公開されている配布パッケージとリポジトリメタデータの間で確認できる条件に留意が必要です。PyPIで公開されている torch_tpu 0.2.0 の配布wheelは、CPython 3.11から3.14（3.14 free-threading版を含む）を対象としたLinux x86-64（glibc 2.31以上）環境向けに提供されています[67, 68]。一方、GitHubの main ブランチに存在する pyproject.toml はソースツリーのバージョンを0.2.1とし、Python 3.11以上、PyTorch 2.11.0以上、libtpu 0.0.50以上を要件として定義しています[69, 70]。ただし、PyPI 0.2.0のパッケージJSONメタデータでは requires_dist および requires_python がnullと設定されているため、main ブランチの下限要件をそのまま0.2.0の確定要件として断定することはできません[70]。また、PyTorch公式の追加プラットフォーム案内ではPython 3.10以上と記載されていますが、実際に配布されているwheelで確認できるのはCPython 3.11以上であり、公開wheelを利用する実務環境ではPython 3.11以上およびglibc 2.31以上を必須条件として扱う必要があります[71, 72]。

インストールの手順において、TorchTPUはランタイムとして libtpu を使用します[73, 74]。リポジトリの案内では、CUDA版PyTorchではなくCPU版PyTorchのwheel index（--extra-index-url https://download.pytorch.org/whl/cpu）を指定してインストールする手順が示されており、既定のCUDA版ビルドを避ける指定が求められます[74, 75]。

クラウドの提供状況と利用コストも事前の確認が必要です。現行の対応世代はTPU v5e、v5p、Trillium（v6e）、Ironwood（v7x）です[66]。Google Cloudのゾーン表では東京リージョン（asia-northeast1-b）にTrilliumが掲載されていますが、同一世代が他のGoogle Cloudサービスでも同様に提供される保証はなく、大規模なチップ構成は利用可能数量が限られる場合があります[76]。そのため、利用地域、世代、サービスごとに現行のゾーンとクォータを直接確認する必要があります[76]。公開価格表における料金単位は「USD/chip-hour」であり、参照時点のオンデマンド料金例として、Ironwood（us-central1）は12.00ドル/chip-hour、Trillium（asia-northeast1 東京）は3.24ドル/chip-hourとなっています[77]。

先行して報道されたx86 CPU向けoneDNN最適化（ZenDNN統合）とTorchTPUでは、技術スタックの階層と導入目的が大きく異なります[78]。oneDNNの取り組みは、AMDのZenDNN最適化をPyTorchのx86 CPU向けoneDNN経路に取り込むものです[78]。これに対しTorchTPUは、まったく新しい専用アクセラレータ（TPU）をPyTorchの外部バックエンドとして追加する統合であり、対象とするハードウェア層が根本から異なります[78]。

実務における選定と検証の軸も明確に分かれます。oneDNNの検証では、PyTorchビルドにおけるMKL-DNN対応の有無を確認した上で対象CPU上での性能を実測することが中心となります（Intelの発表は統合作業中および将来計画を含む段階であり、一律の性能向上を前提とすることはできません）[79]。一方のTorchTPUでは、クラウドインフラにおけるTPUの提供地域、利用サービス、割り当てクォータを確保した上で、モデル内の演算子互換性、動的形状のバケット設定（bounded dynamism）、およびシャーディング構成を網羅的に検証する必要があります[79]。

比較項目	x86 CPU 向け oneDNN（ZenDNN 統合）	アクセラレータ向け TorchTPU	出典
技術スタックの階層	x86 CPU 向け既存バックエンド内の最適化集約	PrivateUse1 ディスパッチャを用いた新規 TPU デバイスバックエンド追加	[78]
実行対象ハードウェア	Intel および AMD の x86-64 CPU	Google TPU（v5e, v5p, Trillium v6e, Ironwood v7x）	[66, 78]
コンパイル・実行経路	oneDNN / MKL-DNN / AOCL-DLP 動的ディスパッチ	Strict eager（DeferredOp） / StableHLO ＋ OpenXLA コンパイラ	[78]
実務上の検証軸	PyTorch ビルドの MKL-DNN 対応確認と対象 CPU 上での性能実測	クラウドの提供ゾーン・クォータ確保、演算子対応、bounded dynamism、分散シャーディング検証	[79]
提供・成熟度ステータス	統合作業中および将来サポート計画の段階（一律の性能向上は前提不可）	バージョン 0.2.0 オープンソース公開済み（PyTorch/XLA は非推奨化）	[67, 79]