Hugging Face ❤️ TorchTPU: TPUs, Natively

TPUs are Google’s custom AI accelerators designed for matrix math in neural networks, used internally for training and serving models like Gemini. Available on Google Cloud from single chips to large pods, they excel at high-compute workloads but historically required TensorFlow or JAX. PyTorch users previously relied on PyTorch/XLA, which introduced complexity such as lazy tensors and difficult debugging. TorchTPU changes this by adding TPU as a native PyTorch backend, allowing standard PyTorch code to run directly on TPUs without workarounds.

TorchTPU integrates with Hugging Face’s ecosystem, enabling native support in Transformers and Diffusers. Models from the Hub can be loaded onto TPUs with a single line of code using device_map="tpu". The backend includes Pallas kernels optimized for TPU hardware, distributed training support, and torch.compile integration for performance gains. Installation on a TPU VM is streamlined via a one-command setup, avoiding unnecessary CUDA dependencies.

Diffusers now treats TPU as a first-class device, supporting both eager and compiled modes for prototyping and deployment. For models too large for a single chip, TensorParallelConfig automatically shards the transformer across multiple TPUs during loading. Many pipelines support TPU out of the box, while others require minimal adjustments via tp_plan. The Kernels Hub provides ready-to-use TPU kernels written in Pallas, offering low-level hardware control and compatibility with Tokamax libraries.

Hugging Face reports significant performance improvements with TPU kernels, such as 21.5× faster decoding for Qwen3-30B-A3B and 12.8× for Gemma 4 26B-A4B at batch size 64 on TPU v6e-8. The company is actively improving TPU support, with 96.5% of test suite passes on TPU and ongoing fixes for failures. Issues discovered during testing have led to upstream bug fixes in Transformers, including StaticCache and Qwen3 tensor-parallel plans.