Hugging Face has introduced native TPU support in Transformers, Diffusers and Kernels through the TorchTPU backend. Developers can now run models from the Hub directly from PyTorch. Transformers and Diffusers no longer require Optimum TPU, which previously provided TPU support for these models.
Standard PyTorch code replaces the XLA route
Previously, running PyTorch on TPUs meant using PyTorch/XLA. Its lazy execution meant operations did not necessarily run when the code invoked them. According to Hugging Face, even small changes could trigger recompilation of the computation graph, and this execution approach made debugging harder.
TorchTPU adds TPUs as another PyTorch backend, allowing developers to use standard PyTorch code. It includes TPU-optimized Pallas kernels, distributed training support and integration with torch.compile.
Transformers v5 loads models directly onto TPUs
In Transformers v5, developers select the target device when loading a model by setting device_map="tpu". Models run in eager mode by default, without requiring compilation. Developers can also use torch.compile or run the model in a distributed setup.
Diffusers supports compilation and splitting across chips
For prototyping or debugging in eager mode, a Diffusers pipeline moves to the TPU with pipe.to("tpu"). Deployment in compiled mode adds a single compilation call while keeping the pipeline code unchanged. There is no need to maintain a separate TPU-specific branch.
A compiled pipeline builds its computation graph on the first call. Subsequent calls reuse that graph as long as the image size stays the same.
For models that exceed a single chip's capacity, developers can pass TensorParallelConfig to from_pretrained. This splits the transformer across chips during loading, so the entire model does not have to fit in one place. Support varies between pipelines, however: some require a tp_plan to be added.
Demonstrations compare compiled and eager execution
Hugging Face demonstrated the difference between the two modes on a Cloud TPU v6e-4. For both demonstrations, Hugging Face split the models across four chips and used bf16 precision and a batch size of 1. MiniMax-H3 generated a 5.2-second video with audio at 1344 × 768 resolution, using 49 steps, in 188 seconds in compiled mode. The same demonstration took 262 seconds in eager mode.
Under the same hardware conditions, Hugging Face demonstrated Qwen-Image-2.1 generating a 1024 × 1024 image in 40 steps. Compiled execution took 3.5 seconds, compared with 8.0 seconds in eager mode.
Kernels Hub offers ready-to-use TPU kernels
The support also includes ready-to-use TPU kernels on Kernels Hub, loaded through the Kernels library. These are written in Pallas, and developers can also integrate kernels from the Tokamax library. One example is ragged-dot for Mixture-of-Experts models.
Hugging Face also reports results comparing a Hub kernel with the default implementation on a TPU v6e-8. At a batch size of 64, the company says decoding was 21.5× faster for Qwen3-30B-A3B and 12.8× faster for Gemma 4 26B-A4B.
Google Cloud access and the state of support
TPUs are available through Google Cloud, ranging from a single chip to pods containing thousands of chips. TorchTPU can be installed on a TPU-equipped virtual machine with one command.
Support is not yet passing every test. Hugging Face says 96.5% of the Transformers tests that run on TPUs pass. The team has reported 34 issues in the TorchTPU repository, of which 19 have already been fixed.
Hugging Face is also working on nightly testing of Transformers on TPUs. It plans to make the results available in the Transformers dashboard.
