Google open-sources TorchTPU, bringing native PyTorch to TPUs

Google has open-sourced TorchTPU, giving developers a native PyTorch backend for training, serving and deploying models on its tensor processing units. Co-developed with Meta’s PyTorch team, the project creates a public software path for targeting Google TPUs through familiar PyTorch APIs.

The release adds TPUs as another native PyTorch target alongside CPUs and GPUs. It follows AMD and Intel’s move to bring PyTorch CPU optimizations into oneDNN, a separate open-source effort to make hardware-specific performance work available through shared software. Google said existing workloads can move to TPUs with minimal code changes, but the evidence reviewed does not establish uniform compatibility or migration effort across all workloads.

Developers can debug models in PyTorch’s eager mode before switching to torch.compile for compiled graph execution. The public TorchTPU repository implements PyTorch ATen kernels for TPUs, includes an XLA compilation cache and supports data-parallel and tensor-parallel distributed training. Licensed under Apache 2.0, the backend translates PyTorch operations into work that Google’s TPU compiler and hardware can execute while retaining the developer-facing PyTorch workflow.

Google is pairing the backend with tools for model training, inference and deployment. Its announcement names vLLM, SGLang, llm-d, TorchTitan, Ray, TorchRec, Hugging Face and LiteRT. The separate vLLM TorchTPU project provides an out-of-tree, PyTorch-native TPU backend that runs standard vLLM model definitions through TorchTPU and torch.compile, alongside TPU-tuned Pallas kernels.

For large-scale training, Google’s TorchTitan distribution provides TPU recipes and TPU-specific training and parallelism paths. Google said TorchTitan has been validated running on as many as 8,000 Ironwood TPU v7x accelerators. It also said a TorchTPU training path through TorchTitan achieved up to 52% model FLOPs utilization on Qwen 3 8B, with near-linear scaling and stability across as many as 8,192 Ironwood TPUs. Google did not publish a full benchmark methodology or independent reproduction in the materials reviewed, leaving those efficiency and scaling figures as vendor-reported results.

Google said Transformers v5 can target TPUs by setting device_map to tpu, while Diffusers supports eager execution, torch.compile and tensor parallelism on TPUs. For edge deployment, the company said LiteRT can convert PyTorch models through LiteRT-Torch. The LiteRT repository documents conversion paths for classic and generative model formats.

The company also introduced a TPU agent suite in early preview. Google described the plug-ins as tools for onboarding, diagnostics, answering user questions and tuning inference performance for PyTorch models.