Hugging Face adds native TPU support via TorchTPU for Transformers and Diffusers
TorchTPU enables PyTorch models from the Hugging Face Hub to run on Google TPUs without torch_xla or JAX, using standard device_map="tpu" calls.
Summary
Hugging Face has integrated TorchTPU, a native PyTorch backend for Google's Tensor Processing Units, across its ecosystem. Transformers v5 models load with device_map="tpu", Diffusers pipelines run via pipe.to("tpu"), and the Kernels Hub offers TPU-optimized Pallas kernels. The company reports 96.5% of Transformers tests pass on TPU and has filed 34 issues against TorchTPU, 19 of which are already fixed. Optimum TPU is no longer required for basic Hugging Face model execution on TPUs.
What changes with native TPU support
- No more torch_xla or JAX required: PyTorch code runs directly on TPU through a native backend, eliminating lazy tensor behavior and graph recompilation issues.
- Single-device abstraction: Models move to TPU with the same
to("tpu")ordevice_map="tpu"pattern used for CUDA. - torch.compile integration: Compiled mode available via
backend="tpu"for optimized execution; Diffusers shows 2-3× speedups in examples. - Distributed loading built-in:
TensorParallelConfigshards models across multiple TPU chips duringfrom_pretrained(), with automatic or customtp_plansupport. - Custom kernels via Pallas: Kernels Hub distributes TPU-optimized kernels; Tokamax ragged-dot kernel achieves 21.5× faster decoding at batch 64 on Qwen3-30B-A3B.
- Installation:
pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpuon TPU VM.
Why it matters
This removes a major friction point for PyTorch users who needed Google Cloud compute but lacked native TPU access, potentially increasing TPU adoption among the large Hugging Face user base. The deprecation of Optimum TPU for basic use simplifies the stack and reduces maintenance surface for both Hugging Face and downstream developers.
Original reference: huggingface.co →