Hugging Face announces native support for TPUs, the acceleration chips designed by Google, within its ecosystem. Thanks to the TorchTPU backend for PyTorch, Transformers, Diffusers, and Kernels run directly on these processors, according to a post published by the company.
A TPU, or Tensor Processing Unit, is a chip designed to accelerate the matrix calculations that occupy most of the time in neural networks. Google has been developing these chips for about a decade and uses them to train and run its own models, including Gemini. They are available on Google Cloud, from a single chip to clusters of several thousand.
Until now, the TPU software ecosystem had been built around TensorFlow and JAX, and PyTorch users had to put in more effort. To run PyTorch on TPU, one had to go through PyTorch/XLA (torch_xla), a solution described in the post as functional but not very convenient: tensors were evaluated lazily, so code didn’t always execute at the expected time, and small changes could trigger graph recompilations. Hugging Face had developed Optimum TPU to offer inference and training of its models on TPU. With TorchTPU, this tool is no longer needed to run Transformers and Diffusers models.
TorchTPU, a Native Backend for PyTorch
TorchTPU adds TPU as a new backend for PyTorch, without modifying the usual PyTorch code. The backend includes Pallas kernels optimized for TPU, support for distributed training, and integration with torch.compile for users seeking more speed. Its installation on a TPU virtual machine is done with a single command: pip install torch-tpu –extra-index-url https://download.pytorch.org/whl/cpu. The post specifies that this index provides the CPU version of torch, which is lighter than the CUDA version that pip downloads by default and which a TPU does not need.
Transformers, Diffusers, and Kernels on TPU
In Transformers v5, whose goal is to be the reference for model definitions so they run everywhere, loading a model onto TPU is done via the device_map=”tpu” parameter. The post illustrates this step with the Qwen/Qwen3.5-9B model. By default, the model runs in eager mode, but it can also use torch.compile or a distributed configuration.
In Diffusers, TPU is now a first-class device. Eager mode, suited for prototyping, requires only a call to pipe.to(“tpu”), while compiled mode, intended for deployment, adds a compilation call. For models too large for a single chip, a TensorParallelConfig passed to from_pretrained distributes the transformer across multiple chips upon loading. The post presents examples run on a Cloud TPU v6e-4, with the MiniMax-H3 models for video and audio, and Qwen-Image-2.1.
The Kernels library provides access to ready-to-use TPU kernels, written in Pallas, which offers low-level control over the hardware. Kernels from the Tokamax library can also be integrated. According to the post, for a batch of 64, decoding is 21.5 times faster on Qwen3-30B-A3B and 12.8 times faster on Gemma 4 26B-A4B, compared to the default implementation, using a TPU v6e-8.
Continuous Testing
Hugging Face indicates it continuously runs the Transformers test suite on TPU and is working on adding nightly continuous integration on these chips, with results to be published on a dedicated dashboard. According to the post, 96.5% of tests run on TPU pass. Each failure results in a reproduction script and a report on the TorchTPU repository: 34 issues have been reported, 19 of which are already fixed. These tests also revealed bugs in Transformers, notably in StaticCache and in Qwen3’s tensor parallelism plans, which were fixed upstream. In case of problems, users are encouraged to open an issue.
