Hugging Face、ネイティブPyTorchバックエンドTorchTPUでTransformersとDiffusersをTPU対応に

従来、PyTorchをTPUで動かすにはPyTorch/XLA（torch_xla）を使う必要があった。テンソルが遅延評価になるため、コードが想定通りのタイミングで実行されない、小さな変更でグラフの再コンパイルが起きる、デバッグしづらいといった癖があったという。Hugging Faceはこれを補うためOptimum TPUを作ったが、TorchTPUがネイティブ対応を提供するため、TransformersとDiffusersでHugging Faceのモデルを動かすのにOptimum TPUはもう不要になるとしている。

TorchTPUはPyTorchにTPUを新しいバックエンドとして加える。通常のPyTorchコードを書くだけでよく、TPU向けに調整されたPallasカーネル、分散学習のサポート、速度を上げたい場合のtorch.compile連携が付属する。Transformersでは`device_map="tpu"`を渡すとモデルがTPUに直接読み込まれ、既定ではeagerモードで動く。例としてQwen/Qwen3.5-9Bのコードが示されている。

DiffusersではTPUがファーストクラスのデバイスになった。試作やデバッグではeagerモードで`pipe.to("tpu")`だけでよく、本番向けにはコンパイル呼び出しを1つ足す。1チップに収まらない大きなモデルは、`TensorParallelConfig`をfrom_pretrainedに渡すとロード時にチップ間で分割される。Cloud TPU v6e-4（TP=4、bf16、バッチ1）の例では、Qwen-Image-2.1の1024×1024、40ステップがコンパイル時3.5秒、eager時8.0秒だった。MiniMax-H3の動画と音声の生成例は、コンパイル時188秒、eager時262秒である。

KernelsのHubにはPallasで書かれたTPU用カーネルがあり、Tokamaxライブラリのカーネルも使える。Tokamaxのragged-dotカーネルをMixture-of-Expertsモデルに組み込むと、バッチが大きいほど効果が増し、バッチサイズ64でQwen3-30B-A3Bが21.5倍、Gemma 4 26B-A4Bが12.8倍高速になった。

品質面では、transformersのテストスイートをTPUで実行しており、TPU上で動くテストの96.5%が通過している。失敗は再現スクリプトにしてTorchTPUリポジトリのissueとして報告し、これまでに34件を報告、19件が修正済みである。TPUでの実行でStaticCacheやQwen3のテンソル並列プランといったTransformers側のバグも見つかり、アップストリームで修正した。TPU向けのnightly CIの追加も進めている。