Google open-sources TorchTPU: PyTorch now runs natively on TPUs

What this is about

Google has released TorchTPU as an open-source project on GitHub (google-pytorch/torch_tpu, Apache 2.0 license). It is a native PyTorch backend for Google's Tensor Processing Units (TPUs), the company's own AI accelerator chips. At the same time, Hugging Face has published a blog post saying the Transformers, Diffusers and Kernels libraries now run directly on TPUs. If you want to use a model from the Hub, you just send it to the "tpu" device.

The problem until now

TPUs are chips built for the matrix math that makes up most of a neural network. Google has been developing them for about ten years and, according to Hugging Face, uses them to train and run its own models, including Gemini. They are available in Google Cloud, from single chips up to pods with thousands of chips.

The obstacle was software. The TPU stack grew up around TensorFlow and JAX. PyTorch users had to go through PyTorch/XLA (torch_xla). That worked, but it had quirks: tensors were evaluated "lazily," so code didn't necessarily run when you expected it to. Small changes could trigger a recompile of the compute graph, and debugging was hard to follow. Hugging Face had therefore built Optimum TPU; according to the blog, TorchTPU makes that unnecessary for Transformers and Diffusers.

What TorchTPU offers

According to the README and the blog, the project includes:

• custom ATen kernels (PyTorch's basic operations) for TPUs
• a compilation cache to run XLA computations faster
• distributed training with data and tensor parallelism (tensor parallelism splits a model across several chips)
• Pallas kernels optimized for TPUs
• integration with torch.compile for graph mode
• examples for Llama, Qwen, ResNet, DLRM, minGPT and UNet, plus tools for debugging and benchmarking

A guide is available at torchtpu.readthedocs.io. The interactive notebooks are Marimo notebooks (similar to Jupyter) and are meant to run on a TPU VM, most easily on a v6e GCE instance. The code itself targets Python 3.11+.

Installation and usage

On a TPU VM, this is all you need:

pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu

The extra index provides the CPU build of torch. Without it, pip downloads the much larger CUDA version, which a TPU doesn't need. Alternatively, you can build the wheel from source with Bazel. The README gives details such as the --config=wheel_common option and --config=no_rbe if you don't have credentials for Google's Remote Build Execution.

Hugging Face integration

In Transformers v5, device_map="tpu" in from_pretrained is enough. The example in the blog loads Qwen/Qwen3.5-9B straight onto the TPU. By default, the model runs in eager mode (one operation at a time); it's faster with torch.compile, and distributed setups are possible too.

In Diffusers, the blog says TPU is a first-class device: for prototypes, pipe.to("tpu") is enough, and for production you add a compile call. The first call compiles; later calls with the same image size reuse the graph. For models that are too large, you pass a TensorParallelConfig, and the model is spread across several chips as it loads.

Examples on a Cloud TPU v6e-4 with TP=4, bf16, batch size 1:

For Qwen-Image-2.1, a 2048×2048 image was also shown, without timing.

Kernels and testing

The Kernels Hub offers ready-made TPU kernels written in Pallas. Kernels from the Tokamax library can also be plugged in. In the example, a ragged-dot kernel is dropped into a Mixture-of-Experts model (Qwen3-30B-A3B-Instruct-2507). At batch size 64, Hugging Face says decoding is 21.5x faster on Qwen3-30B-A3B and 12.8x faster on Gemma 4 26B-A4B (TPU v6e-8, compared with the default implementation). The larger the batch, the bigger the gain.

Hugging Face is running the Transformers test suite on TPU and building a nightly CI. Of the tests that run on TPU, 96.5% pass. Failures are reported as small reproduction scripts and issues in the TorchTPU repo: 34 so far, 19 fixed. The team also found bugs in Transformers itself (StaticCache, Qwen3 tensor-parallel plans) and fixed them.

Our take

TorchTPU lowers the barrier considerably for PyTorch users who want to use TPUs without porting code to JAX. That matters most for teams that need compute in Google Cloud for training or serving and are facing tight capacity on NVIDIA hardware. The speed figures come from Hugging Face itself and compare against the default implementation, not against GPUs; comparisons on cost or against other accelerators are missing. Support is not yet complete: not all tests run, the dependency list in the README is still "TBD," and the version is 0.2.0. To try it out, you need a TPU VM in Google Cloud, which costs money.