NEWS

TorchTPU: Google and Meta Open Native PyTorch Backend for TPUs

News
Reading
5 min
Tables
1
Code
1
In short

New open source project eliminates the code rewrite that was blocking TPU adoption for PyTorch developers, with tests on 5,000 Hugging Face models and benchmarks on TPUv7 chips.

What TorchTPU Is

TorchTPU is an open source project developed by Google in collaboration with Meta that delivers a native PyTorch backend for Google's TPUs (Tensor Processing Units). The announcement came out on November 2, 2026, according to a report by InfoQ. The proposal is to run PyTorch workloads on Google's chips with minimal code modification, supporting eager and compiled execution, distributed computing, and hardware-specific optimizations.

For PyTorch developers, this historically meant choosing between GPU (mature, expensive ecosystem) or TPU (cheap at scale, but full of adoption friction). TorchTPU directly addresses that friction.

Why PyTorch/XLA Wasn't Enough

Google had already been trying to solve this since 2020, via PyTorch/XLA, built on JAX, XLA, and lazy tensors. But Google itself acknowledges that this approach had serious problems.

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.

Google, official TorchTPU announcement

In practice, porting an existing PyTorch project to TPU was laborious: you had to rewrite code around XLA's rules and accept that iterating (training, testing, tuning) was slower than on GPU. For many teams, this migration cost made using TPU unfeasible even when the compute price per hour was lower.

How to Migrate Code from GPU to TPU

TorchTPU flips the priority: eager execution comes first, and existing PyTorch workloads should run with minimal changes. Migrating from GPU to TPU, according to Google, involves switching the device from cuda to tpu and the distributed backend from nccl to tpu_dist.

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()
    world_size = dist.get_world_size()
    ...
    model = model.to(device)

Alternative: use the new torch.accelerator.current_accelerator API to automatically detect which accelerator is available, without hardcoding cuda or tpu in the code. This matters for anyone maintaining a training pipeline that needs to run on both GPU and TPU depending on each provider's availability.

The Three Eager Execution Modes

TorchTPU supports three modes: strict, debug, and DeferAndFuse. The latter is a new hybrid eager/compiled mode that Google claims is 50% or more faster than strict eager mode, while reducing peak HBM memory usage. In practice, DeferAndFuse tries to capture the best of both worlds: the step-by-step debugging convenience of eager mode with part of the performance gain that normally only comes from ahead-of-time compilation.

Testing Against 5,000 Hugging Face Models

To validate that TorchTPU runs existing PyTorch models without modification, Google built an integration pipeline that evaluated around 5,000 Hugging Face models, covering more than 500 architectures, in both eager and compiled mode, across TPU hardware generations starting with v5e.

The model selection methodology combined real market adoption with architectural diversity, avoiding repeated testing of minor variations of the same model and ensuring that production weights were actually evaluated.

The Benchmark Numbers

To measure performance, Google ran Muse Glimmer, a 30-billion-parameter model implemented in standard PyTorch. On a single TPUv7 chip, inference reached 25,700 tokens/s with a 5,120-token sequence and 22,400 tokens/s with a 10,240-token sequence. In training, using four TPUv7 chips, throughput was 5,600 tokens/s and 4,600 tokens/s, respectively, for the same two sequence lengths.

There were also tests with Qwen3 MoE 235B, with training scaling from 1,024 to 8,192 TPUv7 chips. For inference, Google reported competitive performance per dollar on large-scale MoE workloads using vLLM.

ModelHardwareScenarioResult
Muse Glimmer (30B)1x TPUv7Inference, seq. 5,12025,700 tokens/s
Muse Glimmer (30B)1x TPUv7Inference, seq. 10,24022,400 tokens/s
Muse Glimmer (30B)4x TPUv7Training, seq. 5,1205,600 tokens/s
Muse Glimmer (30B)4x TPUv7Training, seq. 10,2404,600 tokens/s
Qwen3 MoE 235B1,024 to 8,192x TPUv7Training at scalenumber not disclosed

What This Means for Brazilian Developers' Wallets

This is the point that matters most for those working outside Silicon Valley: infrastructure costs for training and running models. Top-tier GPUs are scarce and expensive, and in Brazil that's compounded by exchange rates and import taxes when the bill comes down to local hardware. TPU has always been cheaper per hour on Google Cloud, but the barrier of porting PyTorch code there used to eat up much of the savings in engineering time.

With TorchTPU, the migration Google proposes is just a few lines: swap cuda for tpu, nccl for tpu_dist, or use the automatic detection API. This lowers the cost of experimentation: a small team, startup, or academic group in Brazil can run the same checkpoint that already runs on GPU within Google Cloud Platform without rewriting the training pipeline from scratch.

A word of caution: the cost benefit depends on whether the TPU's hourly price in the chosen region offsets the learning curve of a new, still early-stage piece in the stack. TPU availability for accounts in Brazil via Google Cloud is not guaranteed in every region.

The Business Context Behind TPU

Besides selling TPU via Google Cloud Platform, Google also sells the chips to corporate customers to run on-premises, according to Hacker News users who commented on the announcement. Clients cited include Anthropic and, potentially in the future, Meta itself. This adds another layer to the move: Meta, the owner of PyTorch and historically associated with Nvidia infrastructure, helping unlock TPU is also opening up an alternative AI hardware supplier for itself.

Reception in the community was favorable. A user identified as tingkuanpei summed up the general sentiment:

Very interesting work. Google should have done this a long time ago. PyTorch's ecosystem is just so rich.

tingkuanpei, Hacker News user

What's Still Missing

TorchTPU is at an early stage and, according to Google itself, continuously evolving. The public roadmap includes work on dynamic shapes, compilation performance, operator coverage, reinforcement learning, and quantization. For those already using PyTorch in production and considering migrating part of their workload to TPU, this means testing cautiously: not every architecture or custom operator is guaranteed full support today.

The code is available as an open source project, which allows the community, including Brazilian teams, to contribute fixes and operator coverage instead of just waiting on Google's timeline.

Translated from the Brazilian Portuguese original · Read the original