Google and Meta open-source TorchTPU for native PyTorch on TPUs

News update: this post was researched and drafted automatically from the sources listed at the end, and published without a human edit. If you spot a mistake, write to me and I will fix it.

Google and Meta have open-sourced TorchTPU, a native PyTorch backend for Google Tensor Processing Units. The repository sits under google-pytorch on GitHub, the licence is Apache 2.0, and an arXiv paper describes the design as a hybrid eager and compiled architecture. Until now PyTorch users who wanted TPUs had to go through PyTorch/XLA, also called torch_xla. TorchTPU removes that extra layer. You write regular PyTorch and move tensors to 'tpu'.

The package ships with custom ATen kernels, a compilation cache for XLA computations, distributed training support for data and tensor parallelism, and examples for Llama, Qwen, ResNet, DLRM, minGPT and UNet. Hugging Face has already added support in Transformers and Diffusers. A model can be loaded with device_map set to 'tpu' and a pipeline with pipe.to('tpu').

In one line: Google and Meta have shipped TorchTPU, a native PyTorch backend that makes TPUs work like any other PyTorch device with eager execution and torch.compile support.

What shipped

TorchTPU version 0.2.0 is on PyPI for Python 3.11 and above. Installation on a TPU VM is one command:

pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpu

The extra index pulls the CPU build of PyTorch instead of the CUDA build, which matters because TPUs do not need CUDA. The repository uses Bazel for builds and includes custom ATen kernels for PyTorch ops. It has a compilation cache to reuse XLA computations, which should avoid some graph recompilations. Distributed training support covers data parallelism and tensor parallelism. The examples folder includes Llama, Qwen, ResNet, DLRM, minGPT and UNet. There is also a torch.compile integration for graph mode execution.

How this changes PyTorch on TPUs

The old path was PyTorch/XLA. Hugging Face describes the pain: tensors were lazy, code did not always run when you thought it did, small changes could trigger graph recompilations, and debugging was hard. Optimum TPU existed to make that easier. With TorchTPU, Hugging Face says Optimum TPU is not required anymore for Transformers and Diffusers. TPU is now another device.

Transformers v5 accepts device_map='tpu'. Diffusers pipelines accept pipe.to('tpu') for eager mode. For deployment you can add pipe.transformer.compile(backend='tpu', fullgraph=True, dynamic=False). The first call compiles and later calls with the same input shapes reuse the graph. Tensor parallel support is built in through TensorParallelConfig, so models too large for one chip can be sharded as they load.

What Hugging Face has tested

Hugging Face reports that 96.5 percent of the Transformers tests that run on TPU pass. They have reported 34 issues to the TorchTPU repo, and 19 are already fixed. Running on TPU also caught bugs in Transformers itself, including StaticCache and Qwen3 tensor parallel plans. For custom kernels, the Kernels Hub has ready-to-use TPU kernels written in Pallas, and teams can bring in kernels from the Tokamax library. Hugging Face shows a ragged-dot kernel used in a mixture-of-experts model. At batch size 64, decoding is 21.5 times faster on Qwen3-30B-A3B and 12.8 times faster on Gemma 4 26B-A4B compared with the default implementation, on a TPU v6e-8.

What is not sorted yet

The Transformers TPU CI is not all green. The pass rate is high but not complete. Hugging Face says every failure becomes a small repro script and an issue. Some pipelines still need a custom tp_plan for tensor parallelism, though many popular ones work out of the box. The sources here do not give broad benchmark numbers comparing TorchTPU and PyTorch/XLA on the same training jobs. The only specific performance figures are the kernel speedups above. Python support starts at 3.11, and you need a TPU VM to test anything. A v6e GCE instance is the easiest starting point.

What I would try

For our client work, nothing changes immediately unless a project is already on Google Cloud TPUs. Most RAG chatbots and internal AI rollouts run fine on CPU or GPU. But if a client is standardised on Google Cloud and wants lower serving cost, TPU access through native PyTorch is worth watching. I would test TorchTPU on a single v6e VM with a small model like Qwen or Llama before recommending any migration. The install command is simple, but TPU availability and cost still gate the decision. For now, I would keep this on the watchlist and not rewrite working GPU code.