Hugging Face loves TorchTPU for native Tensor Processing Unit support.

Meta description: TorchTPU adds native Tensor Processing Unit support to PyTorch, letting Hugging Face Transformers and Diffusers users run models without Optimum TPU.

Tags: TorchTPU, PyTorch, Tensor Processing Unit, Hugging Face, Diffusers

Featured image: A friendly abstract illustration of a PyTorch-style graph connecting to a chip grid; no real people, no logos, no text.

TorchTPU adds Google's Tensor Processing Unit, or TPU, as a native backend for PyTorch. The change removes the need for Optimum TPU when running Hugging Face models with Transformers and Diffusers. For PyTorch users, it means a Tensor Processing Unit can be treated as another device target rather than a special workflow.

A Tensor Processing Unit is built for matrix math

Google designs Tensor Processing Units as AI accelerators whose main purpose is the matrix arithmetic that neural networks use most. The company has worked on this hardware for about ten years and uses it to train and serve its own models, including Gemini. TPUs are available on Google Cloud, from a single chip up to pods of thousands of chips. They are most useful when a workload needs a lot of compute for training or large-scale serving.

The historical catch was software. TPU tooling matured around TensorFlow and JAX, which made PyTorch adoption less straightforward. Until now, PyTorch users reached TPUs through PyTorch/XLA, known as torch_xla. It worked, but lazy tensors meant code could execute later than expected, small changes could trigger graph recompilations, and debugging was not easy to reason about.

Hugging Face built Optimum TPU to give Hugging Face models inference and training support on TPUs. With TorchTPU bringing native support, Optimum TPU is no longer required to run Hugging Face models with Transformers and Diffusers.

TorchTPU keeps the PyTorch code path simple

TorchTPU registers TPU as a new backend inside PyTorch. Users can continue writing regular PyTorch code. It includes Pallas kernels tuned for TPUs, support for distributed training, and a torch.compile integration for cases that need more speed.

Installing it on a TPU VM takes one command. The extra index matters because it gives the CPU build of torch. Without that index, pip pulls the much larger CUDA build, which is not needed on a TPU.

Transformers and Diffusers can target TPUs directly

Transformers v5 was built around one idea: make model definitions the source of truth so models can run in different places. That includes TPUs now. To load a model on a TPU, users pass device_map="tpu". The model runs in eager mode by default, and torch.compile or a distributed setup can be used when needed.

Diffusers now treats TPU as a first-class device. For prototyping or debugging, eager mode only needs pipe.to("tpu"). For deployment, compiled mode adds one compile call. Existing pipelines and code stay the same, and there is no TPU fork to maintain. The first call compiles the graph, and every later call with the same image size reuses it.

For models too big for one chip, users can pass a TensorParallelConfig to from_pretrained. The transformer is then sharded across chips as it loads, so the full model never has to fit in one place. Many popular pipelines support this out of the box, and the ones that do not are just a tp_plan away. The Diffusers TPU guide has the full recipe. Examples were run on a Cloud TPU v6e-4 with TP=4, each model split across four chips, bf16 precision, and batch 1.

Kernels and tensor parallelism show the performance upside

The Kernels Hub has ready-to-use TPU kernels that can be loaded with the Kernels library. These kernels are written in Pallas, which gives low-level control over the hardware. Kernels from the Tokamax library can also be brought in. The release notes show a ragged-dot kernel from Tokamax plugged into a Mixture-of-Experts model.

The benefit grows with batch size. At batch size 64, decoding is 21.5× faster on Qwen3-30B-A3B and 12.8× faster on Gemma 4 26B-A4B. These decode throughput results were measured on a TPU v6e-8, comparing a Hub kernel with the default implementation.

Testing is still a work in progress

The project is continuously improving TPU support in transformers by running the library test suite. It is working to add a nightly CI on TPU and make the results available in the transformers dashboard at https://huggingface.co/spaces/transformers-community/transformers-ci-dashboard. The Transformers CI dashboard is intended to show TPU results next to other accelerators, but the work is still in progress.

Not every test is green yet. Of the tests that run on TPU, 96.5% pass. Every 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 caught a few bugs in Transformers itself, such as in StaticCache and the Qwen3 tensor-parallel plans. Those were fixed upstream.

What readers can do next

For PyTorch users, the practical effect is that TPUs are no longer a separate, awkward path. The starting point is to install TorchTPU, grab a model from the Hub, and send it to "tpu". The release notes point to TorchTPU on GitHub, TPU kernels on the Hub, and Cloud TPU documentation. If something breaks on a TPU, the project asks users to open an issue so the failure can be found before more users encounter it.