Hugging Face adds native TPU support through TorchTPU PyTorch backend

From torch_xla to a native backend

TPU stands for Tensor Processing Unit. According to Hugging Face, these are Google's own AI accelerators, built for the matrix math that neural networks spend most of their time on. Google has designed them for about ten years and uses them to train and serve its own models, including Gemini. They are available on Google Cloud, from a single chip up to pods of thousands of chips.

The obstacle, in Hugging Face's words, has always been software. The TPU stack grew up around TensorFlow and JAX, so PyTorch users had to work harder to get in. Until now, that meant PyTorch/XLA (torch_xla). It worked, but tensors were lazy, so code did not always run when expected. Small changes could trigger graph recompilations, and debugging was hard to reason about.

Hugging Face had built Optimum TPU to ease this, offering inference and training support for its models on TPUs. With TorchTPU providing native support, the company says Optimum TPU is no longer required to run Hugging Face models with Transformers and Diffusers.

TorchTPU adds TPU as a new backend to PyTorch, so users write regular PyTorch code. It comes with Pallas kernels tuned for TPUs, support for distributed training and a torch.compile integration for more speed. Installation on a TPU VM takes one command: pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu. The extra index supplies the CPU build of torch, since otherwise pip pulls the much bigger CUDA build, which is not needed on a TPU.

Transformers, Diffusers and Kernels on TPU

In Transformers, loading a model on a TPU requires passing device_map="tpu". Hugging Face's example loads Qwen/Qwen3.5-9B this way and generates text. The model runs in eager mode by default, with torch.compile or a distributed setup as options.

In Diffusers, TPU is now a first-class device. Eager mode for prototyping or debugging needs only pipe.to("tpu"), while compiled mode for deployment adds one compile call. Hugging Face says pipelines and code stay the same, with no TPU fork to maintain. The first call compiles, and later calls with the same image size reuse the graph.

For models too big for one chip, a TensorParallelConfig passed to from_pretrained shards the transformer across chips as it loads. Many popular pipelines support this out of the box, and the rest are "just a tp_plan away". On a Cloud TPU v6e-4 with TP=4, bf16 and batch 1, the post lists one example image generated in 3.5 s compiled versus 8.0 s eager (1024×1024, 40 steps) and a video-and-audio clip with MiniMax-H3 at 188 s compiled versus 262 s eager.

The Kernels Hub also hosts ready-to-use TPU kernels written in Pallas, and kernels from the Tokamax library can be used too. With a Tokamax ragged-dot kernel plugged into a Mixture-of-Experts model, Hugging Face reports that at batch size 64 decoding is 21.5× faster on Qwen3-30B-A3B and 12.8× faster on Gemma 4 26B-A4B. These figures compare the Hub kernel with the default implementation on a TPU v6e-8.

Test results and what it means for users

Hugging Face says it is improving TPU support by running the Transformers test suite on TPU, and is working to add a nightly CI on TPU with results on the Transformers dashboard. The company describes this as work in progress: of the tests that run on TPU, 96.5% pass. Each failure becomes a small repro script and an issue on the TorchTPU repository. So far 34 issues have been reported and 19 are already fixed. Running on TPU also exposed bugs in Transformers itself, including in StaticCache and the Qwen3 tensor-parallel plans, which were fixed upstream.

For PyTorch users who want access to Google Cloud TPUs, the practical change is that existing Hugging Face code paths can target the hardware with a device string rather than a separate software stack. Hugging Face's pitch is that "TPUs are no longer a side quest for PyTorch users." The company links to TorchTPU on GitHub, TPU kernels on the Hub and the Cloud TPU documentation, and asks users who hit problems to open an issue.