長年、GoogleのTPUでPyTorchを動かすには、コードをPyTorch/XLAに移植するか、JAXを学ぶしかありませんでした。ネイティブのPyTorchバックエンドであるTorchTPUは、この回り道をなくし、TPUをGPUと同じような、ただのデバイスの1つにします。Hugging Faceによると、Transformers、Diffusers、Kernelsの各ライブラリがこれで直接動くようになり、HubのモデルをTPUに載せるのは設定1つで済みます。
従来の方法には実際のコストがありました。PyTorch/XLAは遅延評価のテンソルを使っていたため、コードが思ったときに実行されるとは限らず、小さな変更で新たなグラフのコンパイルが起き、デバッグも困難でした。Hugging Faceはこれを和らげるためにOptimum TPUを作りましたが、TorchTPUによってTransformersとDiffusersでは不要になりました。新しいバックエンドは既定でeagerモードで動き、速さが必要ならtorch.compileにも対応し、分散学習も扱え、Googleの低水準カーネル言語Pallasで書かれたカーネルを備えています。1つのチップに収まらないモデルは、読み込み時に複数のチップに分割できます。ある例では、TokamaxライブラリのMixture-of-Expertsカーネルにより、バッチサイズ64でのデコードがQwen3-30B-A3Bで21.5倍、Gemma 4 26B-A4Bで12.8倍速くなりました。
最も重要なのはコストです。TPUはGoogle独自のチップで、ニューラルネットワークが時間の大半を費やす行列計算のために作られ、Google Cloudでは1チップから数千チップのポッドまで使えます。PyTorchが一因でNvidiaのGPUにとどまってきたチームは、コードを 書き直さずにTPUの容量を本当の選択肢として見積もれるようになります。AWS Trainium向けの同様のネイティブバックエンドTorchNeuronや、OpenAIのモデルをGoogle Cloudに載せる提携と同時期の発表で、PyTorch、Transformers、Diffusersなど私たちがレビューする道具にとって、Nvidia中心ではない市場の兆しといえます。
完成はしていません。Hugging Faceによると、TPU上で実行されるTransformersのテストの合格率は96.5%で、これまでにTorchTPUのプロジェクトに34件の問題を報告し、毎晩のテストはまだ準備中です。示された高速化は大きなバッチサイズでの特定のカーネルについてのもので、成熟したCUDA環境との全体的な比較ではありません。移行の予算を組む人にとって残る疑問は、自分の作業でTPUの性能と価格がGPUとどう比べられるかで、それは推測ではなく測定で確かめる必要があります。