Google and Meta Launch TorchTPU for Native PyTorch on TPU

TorchTPU is an open-source project developed collaboratively by Google and Meta, providing a native PyTorch backend for Google TPUs. It allows PyTorch workloads to run on TPU hardware with only minor code changes, and supports eager execution, compiled execution, distributed computing, and hardware-specific optimizations.

Google began supporting PyTorch workloads on TPUs in 2020 through PyTorch/XLA, a solution that used JAX, XLA, and lazy tensors. Google says this approach had significant shortcomings:

The key issues are that PyTorch/XLA requires intrusive modifications to source code to explicitly define compilable regions; compilation latency is severe; it does not support PyTorch's native imperative ("eager") execution mode; and it has lagged in supporting emerging distributed computing primitives in the PyTorch ecosystem.

TorchTPU takes a completely different approach, prioritizing eager execution so that existing PyTorch workloads can run with only minor code changes.

Specifically, migrating from GPU to TPU requires replacing the cuda device with tpu and the nccl distributed backend with tpu_dist. Developers can also use the new torch.accelerator.current_accelerator API to identify available accelerators, as shown in the following code snippet.

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)
    ...

Copy code

TorchTPU supports three eager execution modes: Strict, Debug, and DeferAndFuse. DeferAndFuse is a new hybrid eager and compiled execution mode that Google claims "can be more than 50% faster than Strict Eager while reducing peak HBM memory usage."

To test TorchTPU's ability to run existing PyTorch models without modification, Google built an integration pipeline that evaluates approximately 5,000 Hugging Face models across more than 500 architectures in both eager and compiled modes on multiple generations of TPU hardware starting from v5e.

We use a robust testing methodology based on corpus stratification and architecture-balanced selection. This methodology uses a two-stage selection process that balances real-world adoption with broad architectural diversity, avoiding repeated testing of minor variants while ensuring that real production weights are thoroughly evaluated.

To validate TorchTPU's performance, Google ran a 30-billion-parameter Muse Glimmer model implemented with standard PyTorch. On a single TPUv7 chip, inference throughput reached 25,700 tokens per second at a sequence length of 5,120 and 22,400 tokens per second at a sequence length of 10,240. For training, using four TPUv7 chips, throughput reached 5,600 and 4,600 tokens per second at sequence lengths of 5,120 and 10,240, respectively.

Google also conducted training workload experiments with Qwen3 MoE 235B, scaling from 1,024 TPUv7 chips to 8,192. On the inference side, Google says that running large-scale MoE workloads with vLLM can deliver competitive performance per unit cost.

Commenting on the release, Hacker News user tingkuanpei said:

Very interesting work. Google should have done this a long time ago. The PyTorch ecosystem is simply too rich.

Other users pointed out that in addition to offering TPUs through GCP, Google also sells them to "enterprise customers for on-premises deployment", including Anthropic, and that Meta may follow suit in the future.

TorchTPU is still in its early stages and therefore continues to evolve. The project's roadmap includes work on dynamic shapes, compilation performance, operator coverage, reinforcement learning, and quantization. For more details, please refer to the original announcement, which provides a more comprehensive introduction to TorchTPU's architecture and scope.

Original link: https://www.infoq.com/news/2026/11/torchtpu-native-pythorch-tpu/