Which TPU generations does TorchTPU support?

TorchTPU's feature-set page lists TPU v5e, v5p, v6e and v7 as validated. Support for the eighth-generation TPU v8i and v8t chips is marked as near-term work. The getting-started guide uses a v6e Cloud TPU VM as its example and says recent TPU hardware requires GKE.

Does TorchTPU replace PyTorch/XLA?

Google said in its April 2026 TorchTPU announcement that TorchTPU would replace PyTorch/XLA once its public GitHub repository shipped, and that repository is now live. Unlike PyTorch/XLA, which moves tensors to an xla device and traces graphs, TorchTPU uses a native "tpu" device with eager execution and has no JAX runtime dependency.

Which distributed training methods work with TorchTPU?

TorchTPU implements TPU collectives at the torch.distributed process-group level, so it supports DDP, FSDP2, SimpleFSDP and DTensor. The feature page marks checkpointing, activation offloading and multi-slice training up to 8,000 chips as validated, while FSDPv2 and MPMD are experimental and still low-performance.

How much faster is TorchTPU with torch.compile and Pallas kernels?

In TorchTPU's FLUX case study on one TPU v6e chip, a DoubleStreamBlock training step took 178.61 ms in plain eager mode. Adding zero-copy views, fused RoPE, Pallas SplashAttention and torch.compile cut it to 122.51 ms, about 1.46x faster. torch.compile alone gave a 1.12x speedup at that 32,768-token sequence length.

Google · 2026-10-19 · major

TorchTPU goes open source — Google's native PyTorch backend for TPUs

Google has published TorchTPU on GitHub under Apache 2.0, with version 0.2.0 on PyPI. TorchTPU runs PyTorch code on Cloud TPUs by switching the device from "cuda" to "tpu", and it is built to replace PyTorch/XLA.

GitHub card for the google-pytorch/torch_tpu repository

The PyTorch-on-TPU backend Google announced in April is now public code you can pip install.

Quick facts

Maker	Google, with the PyTorch team at Meta
License	Apache-2.0
Version	0.2.0 on PyPI (Linux x86-64, Python 3.11+)
Validated TPUs	v5e, v5p, v6e, v7
Distributed	DDP, FSDP2, SimpleFSDP, DTensor; multi-slice up to 8K chips
Install	pip install torch_tpu

What is it?

Google has opened the TorchTPU repository on GitHub and shipped release 0.2.0 on PyPI, so anyone can now run PyTorch models on Cloud TPUs through it. When Google announced TorchTPU in April 2026, it said the project would replace PyTorch/XLA once the public repository shipped. The repo includes custom ATen kernels, a compilation cache and example models such as Llama, Qwen, ResNet, DLRM, minGPT and UNet.

How does it work?

Porting code mostly means changing device="cuda" to device="tpu" and the process-group backend from "nccl" to "tpu_dist". By default, operators run eagerly one at a time, with an optional DeferAndFuse mode that fuses operator sequences before results come back to the CPU. torch.compile uses a "tpu" Dynamo backend that lowers FX graphs to StableHLO for the OpenXLA compiler, and compiled programs go into a three-tier cache (process memory, host memory and remote storage).

Why does it matter?

PyTorch users who wanted TPU capacity used to go through the PyTorch/XLA compiler bridge or switch to JAX. With an open, native backend, existing PyTorch training scripts can move to TPUs with small edits, and the TPU becomes a more practical alternative to renting NVIDIA GPUs. TorchTPU also has no JAX runtime dependency, so the PyTorch stack stands on its own on TPU hardware.

Who is it for?

ML engineers training or serving PyTorch models on Google Cloud

Frequently asked questions

Which TPU generations does TorchTPU support?: TorchTPU's feature-set page lists TPU v5e, v5p, v6e and v7 as validated. Support for the eighth-generation TPU v8i and v8t chips is marked as near-term work. The getting-started guide uses a v6e Cloud TPU VM as its example and says recent TPU hardware requires GKE.
Does TorchTPU replace PyTorch/XLA?: Google said in its April 2026 TorchTPU announcement that TorchTPU would replace PyTorch/XLA once its public GitHub repository shipped, and that repository is now live. Unlike PyTorch/XLA, which moves tensors to an xla device and traces graphs, TorchTPU uses a native "tpu" device with eager execution and has no JAX runtime dependency.
Which distributed training methods work with TorchTPU?: TorchTPU implements TPU collectives at the torch.distributed process-group level, so it supports DDP, FSDP2, SimpleFSDP and DTensor. The feature page marks checkpointing, activation offloading and multi-slice training up to 8,000 chips as validated, while FSDPv2 and MPMD are experimental and still low-performance.
How much faster is TorchTPU with torch.compile and Pallas kernels?: In TorchTPU's FLUX case study on one TPU v6e chip, a DoubleStreamBlock training step took 178.61 ms in plain eager mode. Adding zero-copy views, fused RoPE, Pallas SplashAttention and torch.compile cut it to 122.51 ms, about 1.46x faster. torch.compile alone gave a 1.12x speedup at that 32,768-token sequence length.

Try it

pip install torch_tpu