Hugging Face Brings Native TPU Support to PyTorch Ecosystem
Hugging Face has introduced native Tensor Processing Unit support within the PyTorch ecosystem, eliminating the need for traditional workarounds like PyTorch/XLA or JAX.
Radwan AbdulhadiEditor and technology writer
Native TPU Support Lands in PyTorch
Running PyTorch workloads on Google's Tensor Processing Units traditionally required porting code through PyTorch/XLA or adopting the JAX framework. According to the [Hugging Face Blog](https://huggingface.co/blog/hf-torchtpu), the introduction of a native PyTorch backend changes this dynamic by treating TPUs simply as another available hardware device.
Tensor Processing Units are specialized custom chips designed by Google specifically for heavy matrix mathematics utilized by neural networks. While these accelerators are widely deployed across Google Cloud infrastructure for training and serving large-scale models like Gemini, software integration challenges historically required extra effort from PyTorch developers.
Simplifying Development Without Optimum TPU
Previously, utilizing TPUs with Hugging Face architectures meant navigating the lazy evaluation behaviors and graph recompilation quirks inherent to PyTorch/XLA. To bridge this gap, Optimum TPU was previously developed to provide inference and training support for Hugging Face models.
With the arrival of native support via [TorchTPU](https://github.com/google-pytorch/torch_tpu), Optimum TPU is no longer necessary to execute Transformers and Diffusers models. Developers can write standard PyTorch code, leverage Pallas kernels tuned specifically for TPU hardware, utilize distributed training setups, and incorporate torch.compile for enhanced performance speed.
Transformers and Diffusers Integration
Transformers v5 has been structured to act as a foundational source of truth for model definitions across diverse hardware environments. Loading a model onto a TPU environment now only requires specifying a device map configuration directly within the loading parameters.
Similarly, the Diffusers library now treats the hardware accelerator as a first-class device. For prototyping workflows, eager mode functions seamlessly with standard pipeline assignments, while compiled execution modes add optimization layers. For models exceeding the capacity of a single chip, developers can distribute transformer layers across multiple chips during the loading process using tensor parallel configurations.
Kernel Optimizations and Ecosystem Testing
The integration also introduces ready-to-use TPU kernels via the Kernels library, providing low-level hardware control through Pallas and external integrations like the Tokamax library. Benchmarking tests on hardware configurations such as Cloud TPU v6e-8 show significant throughput acceleration during decoding tasks for large mixture-of-experts architectures.
To ensure long-term stability, continuous integration testing is expanding to monitor TPU compatibility closely. Out of the active test suite evaluated on TPU infrastructure, 96.5% of tests currently pass successfully, with upstream fixes already applied to address early anomalies discovered during validation.
Sources
- Hugging Face BlogHugging Face ❤️ TorchTPU: TPUs, Natively
Continue chronologically
Related entity coverage
Related stories
TorchNeuron Brings Native PyTorch Support to AWS Trainium
TorchNeuron arrives as a public beta, introducing a native PyTorch backend for AWS Trainium with eager execution and torch.compile support.
Hugging Face Launches Space to Check AWS Trainium Support
Discover which Hugging Face models run on AWS Trainium using the new torch-neuronx model support space, checking compatibility, scripts, and more.
Hugging Face Transformers Now Supports llama.cpp GGUF Models
Hugging Face has added native support for GGUF model files to its Transformers library, enabling developers to run quantized AI locally.
Train Qwen3 on AWS Trainium with torchtitan
Learn how to train Qwen3 on AWS Trainium using torchtitan across three distinct configuration approaches on a single trn2.48xlarge instance.
