Intel / article

Hugging Face ❤️ TorchTPU: TPUs, Natively

Source: huggingface.co
Date: Published Oct 20, 2026

Context

Hugging Face's TorchTPU post says running PyTorch on Google's TPU chips no longer means porting code to PyTorch/XLA or learning JAX. TorchTPU is a native PyTorch backend that adds TPU as an ordinary device, so in Transformers a model loads with device_map set to tpu, and in Diffusers a pipeline moves there with a single .to("tpu") call. The older torch_xla route used lazy tensors, which record operations and run them later, so code did not always execute when it appeared to, small changes could trigger graph recompilation, and debugging was hard to reason about. TorchTPU runs eagerly by default, with torch.compile available for speed, plus distributed training and TPU-tuned Pallas kernels. As a result, Hugging Face says its own adapter, Optimum TPU, is no longer required to run models with Transformers and Diffusers.

The performance figures are Hugging Face's own examples, not an independent . On a Cloud TPU v6e-4 with each model split across four chips, one 1024×1024, 40-step image took 3.5 seconds compiled versus 8.0 seconds eager. Swapping in a ragged-dot kernel from the Kernels Hub for the layers made decoding at batch size 64 21.5 times faster on Qwen3-30B-A3B and 12.8 times faster on Gemma 4 26B-A4B on a TPU v6e-8, compared with the default implementation, and the post notes the gain grows with batch size. Support is not finished: 96.5% of the Transformers tests that run on TPU pass, and of 34 issues the team has reported to TorchTPU, 19 are fixed.

Key takeaways · AI-distilled

TorchTPU is a native PyTorch backend, replacing the torch_xla path whose lazy tensors and graph recompilations made debugging hard; Optimum TPU is no longer needed for Transformers and Diffusers.
Installing with the CPU extra index matters: without it pip pulls the much larger CUDA build of torch, which a TPU VM does not need.
In Transformers, device_map="tpu" loads a model straight onto the TPU in eager mode, with torch.compile and distributed setups available for more speed.
On a Cloud TPU v6e-4 with 4-way tensor parallelism, a 1024x1024 Qwen-Image-2.1 render took 3.5 s compiled versus 8.0 s eager, per the authors' examples.
A Tokamax ragged-dot kernel from the Kernels Hub made batch-64 decoding 21.5x faster than the default on Qwen3-30B-A3B; of Transformers tests run on TPU, 96.5% pass, with 19 of 34 reported issues fixed.

Terms in this piece · Glossary

benchmark — A standard public test set for comparing AI models — the shared scoreboards behind every "model X beats model Y" claim.
mixture-of-experts — A model built from many specialist sub-networks where only a few activate per token, giving big-model capability at small-model running cost.

Read the source huggingface.co