Hugging Face runs Transformers and Diffusers natively on TPUs with TorchTPU

Hugging Face has published a post announcing native TPU support across its libraries, built on TorchTPU, a native PyTorch backend. The pitch: TPU becomes "just another device", so you pick a model from the Hub, send it to "tpu", and you are done. Transformers, Diffusers and Kernels all run natively on TPU.

The post first explains the background. TPUs are Google's own AI accelerators, chips built for the matrix math that neural networks spend most of their time on. Google has been designing them for about ten years and uses them to train and serve its own models, including Gemini. They are available on Google Cloud, from a single chip up to pods of thousands of chips. The catch, the post says, has always been software: the TPU stack grew up around TensorFlow and JAX, so PyTorch users had to work harder to get in.

The old route was PyTorch/XLA (torch_xla). It worked, but tensors were lazy, so code did not always run when you expected, small changes could trigger graph recompilations, and debugging was hard to reason about. Hugging Face built Optimum TPU to ease this, giving its models inference and training support on TPUs. With TorchTPU, the post says, Optimum TPU is no longer required to run Hugging Face models with Transformers and Diffusers.

TorchTPU adds TPU as a new backend to PyTorch, so you write regular PyTorch code. It comes with Pallas kernels tuned for TPUs, distributed training support, and a torch.compile integration for more speed. On a TPU VM it installs with one command: pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu. The extra index gives the CPU build of torch; without it, pip pulls the much bigger CUDA build, which a TPU does not need.

In Transformers, which the post says v5 was built to be the source of truth for model definitions that can run anywhere, you load a model on TPU by passing device_map="tpu" to from_pretrained, and move inputs with .to("tpu"). The example uses Qwen/Qwen3.5-9B. The model runs in eager mode by default; torch.compile or a distributed setup are options.

In Diffusers, TPU is a first-class device. For prototyping or debugging, eager mode needs only pipe.to("tpu"). For deployment, compiled mode adds one compile call, shown as pipe.transformer.compile(backend="tpu", fullgraph=True, dynamic=False). The first call compiles, and later calls with the same image size reuse the graph. Pipelines and code stay the same, with no TPU fork to maintain. For models too big for one chip, you pass a TensorParallelConfig to from_pretrained and the transformer is sharded across chips as it loads. Many popular pipelines support this out of the box, and the others are "just a tp_plan away". The post shows example runs on a Cloud TPU v6e-4 with TP=4 (each model split across four chips), bf16, batch 1; the numbers sit in images.

Kernels: the Kernels Hub has ready-to-use TPU kernels, written in Pallas, and kernels from the Tokamax library can be used too. The example plugs a Tokamax ragged-dot kernel (from the Hub repo tengomucho/ragged-dot-tpu) into a Mixture-of-Experts model, Qwen/Qwen3-30B-A3B-Instruct-2507, via experts_implementation="tokamax_ragged_dot". The benefit grows with batch size. At batch size 64, on a TPU v6e-8, decoding is 21.5× faster on Qwen3-30B-A3B and 12.8× faster on Gemma 4 26B-A4B, compared with the default implementation.

On testing, Hugging Face is running the Transformers test suite on TPU. It is working to add a nightly CI on TPU and show results in the Transformers CI dashboard next to other accelerators; the dashboard is described as work in progress. The post admits it is not all green yet: of the tests that run on TPU, 96.5% pass. Each failure becomes a small repro script and an issue on the TorchTPU repo. So far 34 issues have been reported and 19 are already fixed. Running on TPU also exposed a few bugs in Transformers itself, in StaticCache and the Qwen3 tensor-parallel plans, which were fixed upstream. The post closes by inviting readers to open an issue if something breaks.

Key facts

TorchTPU is a native PyTorch backend that makes TPU a regular device: Transformers uses device_map="tpu", Diffusers uses pipe.to("tpu"), and Kernels loads Pallas TPU kernels from the Hub.
Hugging Face says Optimum TPU is no longer required to run its models with Transformers and Diffusers; the older torch_xla route had lazy tensors, surprise recompilations and hard debugging.
Install on a TPU VM is one command: pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu.
At batch size 64 on a TPU v6e-8, a Hub TPU kernel (Tokamax ragged-dot) makes decoding 21.5× faster on Qwen3-30B-A3B and 12.8× faster on Gemma 4 26B-A4B than the default implementation.
Of the tests that run on TPU, 96.5% pass; 34 issues have been reported on the TorchTPU repo and 19 are fixed, and a nightly TPU CI is still being added.

Why it matters

TPUs are Google's AI chips, available on Google Cloud from one chip to pods of thousands, but the post says the software stack grew up around TensorFlow and JAX, leaving PyTorch users to work harder. Until now the PyTorch route was torch_xla, with lazy tensors, recompilations triggered by small changes and awkward debugging. A native backend means the usual Hugging Face code paths work on TPU with a device string, and the Hub's model definitions are the same ones used elsewhere. The post's own summary: TPUs are "no longer a side quest for PyTorch users".

Who it affects

PyTorch users who wanted TPU capacity on Google Cloud but did not want to port code to torch_xla or learn JAX. Teams that used Optimum TPU, which the post says is no longer needed for Transformers and Diffusers. Diffusers users generating images or video, including those whose models are too large for one chip. People serving Mixture-of-Experts models, who are the ones shown the Hub kernel speedups.

How to use it

On a TPU VM, run pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu; the extra index avoids pulling the larger CUDA build of torch. In Transformers, call AutoModelForCausalLM.from_pretrained with device_map="tpu" and move tokenized inputs with .to("tpu"); the model runs in eager mode, and torch.compile or distributed setups are optional. In Diffusers, call pipe.to("tpu") for eager mode, and add a compile call for deployment; pass a TensorParallelConfig to from_pretrained to shard a large model across chips. For kernels, use get_kernel from the Kernels library and register it, as in the ragged-dot example for a Mixture-of-Experts model. The Diffusers TPU guide has the full recipe. The post gives no pricing for Cloud TPUs.

How solid is it

This is a first-party Hugging Face post, so the claims are the company's own. The 21.5× and 12.8× figures are relative speedups at batch size 64 on a TPU v6e-8, measured against "the default implementation"; the post gives no absolute throughput in text, since the benchmark charts are images. The 96.5% figure covers only the tests that run on TPU; the total number of tests or how many are skipped is not stated. The post names no authors, release date or TorchTPU version.

Risks and caveats

Support is not complete: the post says it is "not all green yet", and 34 reported issues include 15 not yet fixed (19 of 34 are fixed). A nightly CI on TPU is something the team is working to add, not something running today, and the dashboard is marked work in progress. The headline speedups come from one kernel on two Mixture-of-Experts models at one batch size, and the post says the gain grows with batch size, so smaller batches should not be assumed to match. The Kernels example loads a kernel from a Hub repo with trust_remote_code set for that repo, so readers copying it are running code from that repo. No comparison of TPU against GPU speed or cost is given.

“TPUs are no longer a side quest for PyTorch users.”

— Hugging Face blog post, "Hugging Face ❤️ TorchTPU: TPUs, Natively"