Hugging Face TorchTPU Brings Native PyTorch to Google TPUs
Hugging Face brings native Google TPU support to PyTorch through TorchTPU, enabling Transformers, Diffusers, distributed training and faster Hub kernels.
Summary
Hugging Face introduced TorchTPU on October 20, 2026, making Google TPUs native PyTorch devices instead of requiring torch_xla, JAX or Optimum TPU. Google has developed TPUs for about 10 years, uses them for Gemini, and offers configurations from one chip to pods with thousands on Google Cloud. TorchTPU supports regular PyTorch, eager execution, Pallas kernels, distributed training and torch.compile. Installation uses `pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu`, avoiding an unnecessary CUDA build.
Transformers v5 loads Hub models such as Qwen/Qwen3.5-9B with `device_map="tpu"`. Diffusers moves pipelines including Stable Diffusion 3.5 Medium with `pipe.to("tpu")`; compiled graphs are reused for matching image sizes. TensorParallelConfig shards oversized models across chips while loading. On Cloud TPU v6e-4 with four-way tensor parallelism, bf16 and batch one, MiniMax-H3 generated 1344 by 768 video and audio in 188 seconds compiled versus 262 eager, while Qwen-Image-2.1 produced a 1024 by 1024 image in 3.5 seconds versus 8.0; a 2048 by 2048 example also ran. On TPU v6e-8 at batch 64, a Tokamax ragged-dot Pallas kernel accelerated decoding 21.5 times for Qwen3-30B-A3B and 12.8 times for Gemma 4 26B-A4B. Of applicable Transformers tests, 96.5% pass; 34 TorchTPU issues were filed, 19 are fixed, and TPU testing exposed upstream StaticCache and Qwen3 tensor-parallel bugs. Nightly TPU CI and public dashboard results remain in development.
Positives
- TorchTPU makes Google TPUs native PyTorch devices, removing the need to port Hugging Face workloads through torch_xla, JAX or Optimum TPU.
- Transformers v5 can place Hub models directly on TPUs using `device_map="tpu"`, with eager, compiled and distributed execution available.
- Diffusers supports unchanged TPU pipelines, while TensorParallelConfig shards models across multiple chips before the complete model must occupy one device.
- A Tokamax Pallas kernel delivered 21.5 times faster Qwen3-30B-A3B decoding and 12.8 times faster Gemma 4 26B-A4B decoding at batch 64.
- 96.5% of applicable Transformers tests pass on TPUs, and 19 of 34 reported TorchTPU issues have already been fixed.
Risks & concerns
- 3.5% of applicable Transformers tests still fail, so native TPU compatibility is not yet complete.
- Nightly TPU continuous integration and public results in the Transformers dashboard remain unfinished.
- Compiled Diffusers execution incurs compilation on the first call, and graph reuse requires subsequent calls with the same image size.
- Installing without the specified CPU PyTorch index causes pip to download the much larger, unnecessary CUDA build.