Google and Meta Open-source TorchTPU to Bring Native PyTorch Support to Google TPUs

Google, working with Meta, released TorchTPU, an open-source native PyTorch backend for Google TPUs that lets existing PyTorch workloads run on TPU hardware with minimal code changes, such as swapping the `cuda` device for `tpu` and the `nccl` distributed backend for `tpu_dist`. Google says its DeferAndFuse hybrid eager/compiled mode "can be 50% or more faster than Strict Eager while reducing peak HBM memory," and it validated the backend on about 5,000 Hugging Face models across 500+ architectures on TPU generations from v5e, measuring 25.7k token/s inference at 5120 sequence length on a single TPUv7 chip for a 30-billion parameter Muse Glimmer model. The release replaces the 2020 PyTorch/XLA approach, which Google says required intrusive source-code changes, suffered severe compile-time latency and lacked native eager execution, and TorchTPU remains at an early stage with a roadmap still in progress.

TorchTPU is an open-source project developed by Google in collaboration with Meta that provides a native PyTorch backend for Google TPUs https://pytorch.org/blog/introducing-torchtpu/ . It enables PyTorch workloads to execute on TPU hardware https://cloud.google.com/tpu with minimal code modifications, supporting eager and compiled execution, distributed computing, and hardware-specific optimizations. Google began enabling PyTorch workloads to run on TPUs in 2020 through PyTorch/XLA, which leveraged JAX, XLA https://github.com/openxla/xla , and lazy tensors. According to Google, this approach had significant shortcomings: Crucially, PyTorch/XLA required intrusive source-code modifications to explicitly define compilable regions, suffered from severe compile-time latency, lacked support for PyTorch's native imperative “eager” execution mode, and lagged in supporting emerging distributed computing primitives within the PyTorch ecosystem. TorchTPU https://github.com/google-pytorch/torch tpu takes a fundamentally different approach, making eager execution a top priority, and enabling existing PyTorch workloads to run out with minor code modifications. Specifically, migrating from GPUs to TPUs involves replacing the cuda device with tpu and the nccl distributed backend with tpu dist . Alternatively, developers can use the new torch.accelerator.current accelerator API to identify the available accelerator, as shown in the following snippet. python import torch import torch.distributed as dist def main : device = torch.device "cuda" - device = torch.accelerator.current accelerator dist.init process group backend="nccl" - dist.init process group rank = dist.get rank word size = dist.get world size ... model = model.to device ... TorchTPU support three eager execution modes: strict, debug, and DeferAndFuse. The latter is a new hybrid eager/compiled execution mode that Google claims "can be 50% or more faster than Strict Eager while reducing peak HBM memory". To test TorchTPU's ability to run existing PyTorch models with no modifications, Google built an integration pipeline to evaluate about 5,000 Hugging Face models across 500+ architectures in both eager and compiled modes on TPU hardware generations starting from v5e. We use a robust testing methodology grounded in corpus stratification and balanced architectural selection. Our approach relies on a two-tiered selection process designed to balance real-world adoption with broad architectural diversity, avoiding redundant testing of minor variations while ensuring real production weights are thoroughly evaluated. To validate TorchTPU's performance, Google ran a 30-billion parameter Muse Glimmer model implemented in standard PyTorch. On a single TPUv7 chip, inference throughput reached 25.7k token/s at 5120 sequence length and 22.4k token/s at 10,240 sequence length. For training, using four TPUv7 chips, throughput reached 5.6k token/s and 4.6k token/s at 5120 and 10,240 sequence length, respectively. Google also conducted experiments using Qwen3 MoE 235B for training workloads, scaling from 1,024 to 8,192 TPUv7 chips. For inference, Google reported competitive performance-per-dollar for large-scale MoE workloads using vLLM. Commenting on the announcement, Hacker News user tingkuanpei observed https://news.ycombinator.com/item?id=50180165 : Very interesting work. Google should have done this a long time ago. PyTorch's ecosystem is just so rich. Other users pointed out that beyond offering TPUs through GCP, Google sells them to "enterprise client for on premises deployments" https://news.ycombinator.com/item?id=50175919 , which include Anthropic https://news.ycombinator.com/item?id=50171644 and in Meta following potentially in the future https://news.ycombinator.com/item?id=50174910 . TorchTPU is still at an early stage and therefore in continuous evolution. The project's roadmap include work on dynamic shapes, compilation performance, operator coverage, reinforcement learning, and quantization. For further details, make sure to read the original announcement, which provides a more comprehensive view of TorchTPU's architecture and scope.