Google opened TorchTPU, a native PyTorch backend for TPUs

Writing PyTorch code for Google TPUs used to mean porting it to torch_xla or learning JAX. On 20 October, Google released TorchTPU as open source. It adds a native tpu device to PyTorch, so sending a model to the TPU reads the same as sending it to a GPU.

The announcement reached developers through two posts the same day: Google's own cloud blog and a longer walkthrough from Hugging Face, written with Google engineers. The Hugging Face post is the detailed one. It describes a working integration with Transformers and Diffusers, lists which models run and which still fail, and gives measured numbers rather than targets.

What the integration looks like

TorchTPU installs with a single pip command and exposes TPUs as a regular device. Transformers reads device_map="tpu" the way it reads device_map="cuda". Pallas kernels, Google's low-level TPU library, provide tuned matrix operations. The library also integrates torch.compile, Google's own torch_xla interoperability layer, and distributed training through PyTorch's standard DeviceMesh and TensorParallelConfig.

Hugging Face said it ran its Transformers test suite on TorchTPU and 96.5 percent of the tests passed. The team filed 34 issues against the library during the integration, and 19 were fixed before the release.

The numbers

Google and Hugging Face measured custom Pallas kernels against a baseline PyTorch implementation on a Cloud TPU v6e-8 at batch size 64. A ragged-dot kernel built for mixture-of-experts models decoded 21.5 times faster on Qwen3-30B-A3B and 12.8 times faster on Gemma 4 26B-A4B. Both are vendor measurements from the authors on their own hardware.

The compile step hurts first-token latency on short runs. The post gives a video generation example at 5.2 seconds compiled against 8.0 seconds for the eager mode a GPU user would recognise, both in a single image step.

Which models already work

Transformers can run Qwen3.5-9B, Qwen3-30B-A3B, Gemma 4 26B-A4B and several Mixture-of-Experts models on TPU through the native path. Diffusers has running recipes for Stable Diffusion 3.5, MiniMax-H3 and Qwen-Image-2.1. The post lists the models currently blocked, which include some video and multimodal pipelines.

TorchTPU retires the Optimum TPU adapter, which Hugging Face used to maintain for TPU users and which carried its own version drift. A team that had wired its training to Optimum TPU can read a migration note in the post.

What a team actually gets

The library runs on Cloud TPU VMs: v6e-4 and v6e-8 for the tests, and the same path scales to TPU pods of thousands of chips through Google Cloud. A PyTorch team that chose GPUs because the CUDA path was the only one its code worked on now has a second hardware vendor it can bid against the first. For a team already running on GPUs today, the practical question is whether the per-step cost of a TPU plus the compile step beats their current cost per step on an Nvidia card. The post does not answer that; a team has to measure its own job.

Source

Hugging Face ❤️ TorchTPU: TPUs, Natively and TorchTPU is now open source: run native PyTorch on TPUs by Google Cloud.