Hugging Face、PyTorch向けTPUバックエンド「TorchTPU」を公開

Hugging Faceは、GoogleのAIアクセラレータであるTPUをPyTorchから直接利用可能にするバックエンド「TorchTPU」を公開した。これにより、TransformersやDiffusersなどのライブラリがTPU上でネイティブ動作するようになる。

これまでPyTorchでTPUを利用するには「torch_xla」を介する必要があり、グラフの再コンパイルやデバッグの複雑さが課題だった。同社は以前「Optimum TPU」を提供していたが、TorchTPUの登場により、今後はこれを介さず実行可能になる。

TorchTPUはPyTorchのバックエンドとして機能し、Pallasカーネルや分散学習、torch.compileをサポートする。インストールは「pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu」で実行できる。

Transformers v5では、モデル読み込み時にdevice_map="tpu"を指定するだけで、モデルをTPU上で直接ロード可能だ。Diffusersでもパイプラインに対してpipe.to("tpu")を実行すればTPUを利用できる。TensorParallelConfigを使えば、単一チップに収まらない大規模モデルを複数チップに分割してロード可能だ。

Kernelsライブラリを使えば、TPU向けに最適化されたカーネルをモデルに適用できる。例えばMixture-of-ExpertsモデルでTokamaxライブラリのカーネルを利用した場合、バッチサイズ64においてQwen3-30B-A3Bで21.5倍、Gemma 4 26B-A4Bで12.8倍のデコード速度向上が確認された（TPU v6e-8環境での測定）。

Hugging FaceはTransformersライブラリのテストスイートをTPU上で継続的に実行しており、現時点でテストの96.5%が通過している。発見された不具合はTorchTPUリポジトリで管理され、これまでに34件の課題を報告し、19件が修正済みだ。なお、TPUの仕様や詳細な技術要件については言及されていない。

開発者にとっての意味

既存のPyTorch環境でTPUを利用する開発者は、即座に移行すべきではない。TorchTPUはOptimum TPUを介さず直接TPUを利用可能にするが、テスト通過率は96.5%に留まり、未解決の不具合も残る。まずは手元のワークロードで、既存のtorch_xla環境とTorchTPU環境の安定性を比較検証し、デバッグの容易さを確認する必要がある。記事内のデコード速度向上は特定のモデルと環境での測定であり、測り方は原文では触れていない。自社の推論パイプラインで同等の速度が出るか、コスト対効果を実測するまで、本番環境のコードは変えない。