Open AI Models

Hugging Face adds native TPU support via TorchTPU for Transformers and Diffusers

Key Takeaway

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") or device_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: TensorParallelConfig shards models across multiple TPU chips during from_pretrained(), with automatic or custom tp_plan support.
  • 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/cpu on 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 →