- TorchTPUはTPUをPyTorchの新バックエンドとして追加し、通常のPyTorchコードで実行できる
- Transformersはdevice_map="tpu"、Diffusersはpipe.to("tpu")の指定でTPU上で動く
- TPU上のテストは96.5%が通過、34件の不具合を報告し19件は修正済み
TPUをPyTorchのバックエンドに
Hugging Faceは2026年10月20日、GoogleのAIアクセラレータTPUをPyTorchから直接利用できるようにする「TorchTPU」への対応を発表しました。TPUはGoogleが約10年設計してきたチップで、ニューラルネットワークが多くの時間を費やす行列演算に特化しています。Google Cloudでは1チップから数千チップのポッドまで提供され、Geminiの学習や提供にも使われています。一方でTPUのソフトウェア資産はTensorFlowとJAXを中心に育っており、PyTorchユーザーは追加の手間を強いられてきました。
TorchTPUはTPUをPyTorchの新しいバックエンドとして追加するもので、利用者は通常のPyTorchコードを書けばよいとされています。TPU向けに調整されたPallasカーネル、分散学習のサポート、速度を追求したいときのためのtorch.compile連携が含まれます。TPU VM上ではpip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu の1コマンドで導入できます。余分なインデックス指定は、TPUでは不要なCUDA版ではなくCPU版のtorchを取得するためのものです。
TransformersとDiffusersの対応
Transformers v5はモデル定義の「正となる情報源」となり、どこでも動かせることを狙って作られました。その「どこでも」にTPUが加わった形で、モデルをTPUに載せるにはfrom_pretrainedでdevice_map="tpu"を渡すだけで済みます。既定ではeagerモードで動作し、必要に応じてtorch.compileや分散構成を選べます。記事ではQwen/Qwen3.5-9Bを例に、推論コードがそのまま動くことが示されています。
画像・動画生成のDiffusersでは、TPUが一級のデバイスとして扱われます。試作やデバッグならpipe.to("tpu")だけでeagerモードで動き、配備向けにはコンパイル用の呼び出しを1つ足します。初回呼び出しでコンパイルされ、同じ画像サイズの以降の呼び出しではグラフが再利用されます。1チップに収まらないモデルは、from_pretrainedにTensorParallelConfigを渡すことで読み込みながらチップ間に分割できます。記事ではstabilityai/stable-diffusion-3.5-mediumを使った例が示されています。
カーネルと性能の数値
Kernels HubにはTPU向けの既製カーネルが用意されています。TPUカーネルはPallasで書かれ、ハードウェアを低レベルに制御できます。Tokamaxライブラリのカーネルも取り込め、記事ではragged-dotカーネルをMoE(Mixture of Experts)モデルに組み込む例が示されています。
バッチサイズを大きくするほど効果は大きく、batch size 64ではデコードがQwen3-30B-A3Bで21.5倍、Gemma 4 26B-A4Bで12.8倍高速になったと報告されています。これはTPU v6e-8上でHubのカーネルと既定実装を比べた数値です。Diffusersの例はCloud TPU v6e-4上で、TP=4(各モデルを4チップに分割)、bf16、batch 1という条件で実行されています。
まだ緑ではないテスト
Hugging FaceはTransformersのテストスイートをTPU上で回し続けており、TPUで実行されるテストの96.5%が通過しているといいます。失敗は再現スクリプトとTorchTPUリポジトリのissueになり、これまでに34件を報告、うち19件は修正済みとされています。TPU上での実行を通じて、StaticCacheやQwen3のテンソル並列プランなどTransformers本体のバグも見つかり、アップストリームで修正されました。
夜間CIをTPU上で動かし、結果をTransformersのCIダッシュボードで公開する作業も進めています。ただし現時点ですべてが成功しているわけではなく、記事も「まだ全部が緑ではない」と認めています。TPUの挙動がCUDA環境と同一とは限らないため、導入時には自分のモデルで検証する必要があると考えられます。
従来方式との違い
これまでPyTorchでTPUを使うにはPyTorch/XLA(torch_xla)を経由する必要がありました。動作はしたものの、テンソルが遅延評価されるためコードが思ったタイミングで実行されず、小さな変更がグラフの再コンパイルを招き、デバッグもしにくいという課題がありました。TPUのソフトウェア資産がTensorFlowとJAXを中心に育ってきたことも、PyTorchユーザーにとっての障壁でした。
Hugging Faceはこの隙間を埋めるためにOptimum TPUを用意し、Hugging Faceのモデルに対するTPU上の推論と学習を支えてきました。TorchTPUのネイティブ対応により、TransformersとDiffusersでHugging Faceモデルを動かす限りOptimum TPUは必須ではなくなりました。記事はTPUがPyTorchユーザーにとっての「寄り道」ではなくなったと表現しています。
| 項目 | PyTorch/XLA(torch_xla) | TorchTPU |
|---|---|---|
| 実行モデル | テンソルが遅延評価され、書いた通りに走るとは限らない | 既定はeager、速度が必要ならtorch.compile |
| 再コンパイル | 小さな変更でグラフ再コンパイルが起きうる | コンパイルは任意で、使う場合のみ |
| デバッグ | 挙動の把握が難しい | 通常のPyTorchコードとして記述 |
| 導入 | torch_xla経由(Optimum TPUも併用) | pip install torch-tpuの1コマンド |
TPUs are no longer a side quest for PyTorch users. Install TorchTPU, grab a model from the Hub, and send it to "tpu".
TPUはもはやPyTorchユーザーにとっての寄り道ではありません。TorchTPUをインストールし、Hubからモデルを取得して「tpu」に送るだけです。
今後の見通し
TorchTPUは発表されたばかりで、TPU上のテスト通過率は96.5%、報告済み34件のうち19件が修正済みという段階です。今後は夜間CIの整備が進み、TransformersのCIダッシュボードでTPUの結果が他アクセラレータと並べて確認できるようになるとみられます。Optimum TPUを使わずに済む範囲がどこまで広がるか、またPallasで書かれたカーネルがKernels Hubにどれだけ集まるかが、実務での採用判断を左右しそうです。Cloud TPUの世代ごとの性能差や、日本国内リージョンでの提供状況が明らかになれば、より具体的な判断ができます。
日本の開発者・IT企業にとっての意味
Google Cloudを基盤にしている日本のIT企業にとって、PyTorchの資産をそのままTPUへ持ち込める意味は小さくありません。これまではTensorFlowやJAX、あるいはPyTorch/XLAの癖を理解した上でTPUを選ぶ必要があり、モデル定義を二重に持つような設計になりがちでした。Transformersでdevice_map="tpu"、Diffusersでpipe.to("tpu")という指定で済むなら、GPU向けに書いたコードの差分を小さく保ったまま、推論や学習の計算資源の選択肢にTPUを加えられます。特に大規模なモデルを複数チップに分割して載せたい場合、TensorParallelConfigやtp_planの仕組みは検討材料になります。ただしTPUが使えるかどうかは契約とリージョン次第であり、性能の数値も特定のTPU世代と条件に限ったものです。まずは小規模なパイプラインで検証し、コストと速度を見極める進め方が現実的といえます。
気になる点
- 日本国内のGoogle CloudリージョンでもTorchTPUは使えますか
- 原文はGoogle CloudのTPUとTPU VM上での手順を示すのみで、リージョンごとの提供状況には言及していません。TPUが利用できる環境であれば、pip install torch-tpuで導入しdevice_map="tpu"を指定する流れは同じとみられます。国内リージョンでの提供有無はCloud TPUの公式情報で確認する必要があります。
- Optimum TPUを使っている既存コードは書き換えが必要ですか
- TransformersとDiffusersでHugging Faceモデルを動かす限り、TorchTPUのネイティブ対応によりOptimum TPUは不要になったと記事は説明しています。Transformersではdevice_map="tpu"、Diffusersではpipe.to("tpu")を指定する形に置き換えることになります。それ以外のライブラリや独自実装については原文に記述がありません。
- TPUの料金はどうなりますか
- 記事にはTPUの料金に関する記述がなく、現時点では公表されていません。TPUはGoogle Cloud上で1チップから数千チップのポッドまで提供されると説明されているため、費用はGoogle CloudのTPU利用料金に従うとみられます。具体的な金額はCloud TPUの料金ページで確認が必要です。
用語解説
- TPU
- Googleが設計するAI向けアクセラレータ。行列演算に特化し、Google Cloudで1チップから数千チップのポッドまで利用できる。
- PyTorch/XLA(torch_xla)
- PyTorchのプログラムをXLAコンパイラ経由でTPUなどのアクセラレータ上で動かすための従来の仕組み。
- Pallas
- TPUなどのハードウェアを低レベルに制御してカーネルを書くための仕組み。TPU向けカーネルの記述に使われる。
- torch.compile
- PyTorchのモデルをコンパイルして高速化する機能。TorchTPUでは任意で利用し、速度を優先するときに使う。
- eagerモード
- 計算グラフを組まずに演算をその場で実行する方式。挙動が分かりやすくデバッグしやすいが最適化は限られる。
- テンソル並列
- 1つのモデルを複数のチップに分割して載せる手法。DiffusersではTensorParallelConfigで指定する。
出典
Hugging Face ❤️ TorchTPU: TPUs, Natively
AI導入も顧問も、お任せください
何から手をつけるか、どこまでAIに任せるか。自社の業務に合わせて整理し、導入から運用まで伴走します。相談だけでも構いません。
AI導入について相談する