Google and Meta have developed TorchTPU, an open-source PyTorch backend for Google TPUs. It is designed to run existing PyTorch workloads with relatively small code changes, while supporting eager execution, compiled execution and distributed computing.
Google has supported PyTorch on TPUs through PyTorch/XLA since 2020. According to Google, that approach can require intrusive changes to define compilable regions, has substantial compilation latency, does not support PyTorch’s native imperative eager-execution model, and has lagged behind newer distributed-computing features in the PyTorch ecosystem.
For workloads moving from GPUs, the changes described for TorchTPU include replacing the cuda device with tpu and the nccl distributed backend with tpu_dist. Developers can also use torch.accelerator.current_accelerator to identify an available accelerator.
TorchTPU offers three eager-execution modes: Strict, Debug and DeferAndFuse. The last combines eager and compiled execution. Google claims it can run more than 50% faster than Strict Eager while reducing peak high-bandwidth-memory use; that figure is a vendor claim, not an independent benchmark.
Google says it built an integration pipeline to evaluate roughly 5,000 Hugging Face models spanning more than 500 architectures, in eager and compiled modes, across multiple TPU generations starting with v5e. In a separate test using a standard PyTorch implementation of the 30-billion-parameter Muse Glimmer model, Google reported inference throughput on one TPUv7 chip of 25,700 tokens per second at a sequence length of 5,120 and 22,400 at 10,240. For training on four TPUv7 chips, it reported 5,600 and 4,600 tokens per second at those respective sequence lengths. Google also tested a Qwen3 MoE 235B training workload at scales from 1,024 to 8,192 TPUv7 chips.
TorchTPU remains at an early stage. Its roadmap includes work on dynamic shapes, compilation performance, operator coverage, reinforcement learning and quantization.