Google Open-Sources TorchTPU, Making TPU a Native PyTorch Device

4,917 Models, 97% GPU Alignment: Google Turns TPU into a Native PyTorch Device

If this news could be summarized in a single line of code, it would be device_map="tpu".

At the PyTorch Conference 2026 held in San Jose, Google announced the open-sourcing of TorchTPU, a complete software stack that enables PyTorch to run natively on TPUs. From now on, deploying models to TPUs using Hugging Face Transformers v5 requires passing just this one parameter.

Behind this change of a single parameter lies a strategic pivot in Google's software roadmap, and represents the most direct challenge yet to NVIDIA's CUDA ecosystem in nearly a decade.

Competition in AI chips has never been solely about the hardware itself. PyTorch is the world's most widely used AI framework, and since its inception in 2016, it has been deeply coupled with CUDA. NVIDIA engineers have long invested heavily to ensure PyTorch performs optimally on their GPUs; Wall Street analysts even view CUDA as NVIDIA's strongest shield against competitors. In contrast, Google's TPUs have primarily served TensorFlow and JAX for the past decade. For PyTorch users wanting to utilize TPUs, they had to go through an "adapter" layer like PyTorch/XLA, incurring significant additional engineering costs.

TorchTPU changes this premise entirely. According to Google's official statement, TPUs are now native options in PyTorch, standing alongside GPUs and CPUs.

• Project Repository: https://github.com/google-pytorch/torch_tpu

The project was jointly developed by Google and Meta (the parent organization of PyTorch). Major open-source projects such as vLLM, SGLang, Hugging Face, Ray, and TorchTitan have simultaneously completed integration. A 23-page technical report was also released on arXiv at the same time.

• Report Title: TorchTPU – A Hybrid Eager/Compiled Architecture for Native PyTorch Execution on TPU
• Report Link: https://arxiv.org/pdf/2610.24161

The timing is also noteworthy. Over the past year, Anthropic signed agreements for up to 1 million TPUs, and Meta was reported to have reached multi-billion dollar contracts with Google for TPU rentals. The hardware is selling well; software was the last remaining bottleneck. Now, Google is stepping in to close that gap.

Not Just Another "Adapter"

Running PyTorch on TPUs is not new. The technical report notes that Google launched PyTorch/XLA based on lazy tensors as early as 2020, candidly acknowledging its issues: it required intrusive source code modifications to define compilable regions, suffered from severe compilation latency, lacked support for PyTorch's native eager execution mode, and lagged in adopting new distributed primitives. Later, TorchAX translated PyTorch operators individually into JAX executions. From the report's perspective, these solutions essentially forced PyTorch to accommodate JAX's execution model.

When TorchTPU was initiated in mid-2025, several core principles were established: no runtime dependency on JAX, robust eager execution, and making the migration of code from GPU to TPU as simple as changing .to("cuda") to .to("tpu").

In terms of implementation, it is built upon PyTorch's PrivateUse1 interface reserved for third-party hardware backends. It registers as a new C++ backend within the PyTorch dispatcher, providing developers with standard PyTorch tensors—without subclasses or wrappers.

Comparison of device initialization code before and after migrating from GPU to TPU.

Eager execution is implemented in three tiers:

• The default Strict mode dispatches operators individually and executes them asynchronously with the CPU, mirroring the standard experience on GPUs.
• Debug mode executes synchronously, useful for troubleshooting shape errors and NaNs.
• The third tier, DeferAndFuse, is a new hybrid eager-compilation mode introduced by TorchTPU. It batches a sequence of operators and automatically fuses them into larger computation blocks sent to the TPU only when values are needed. This results in speeds over 50% faster than Strict mode while also reducing peak memory usage.

All three modes share a hierarchical compilation cache, allowing identical operator sequences to be compiled once and reused repeatedly.

For maximum performance, developers can switch to torch.compile. Here, Google made a deliberate trade-off: after capturing the computation graph using TorchDynamo, it bypasses PyTorch’s default Inductor compiler. Instead, it maps operators to StableHLO and hands them off to XLA to generate TPU programs.

The report provides three reasons for this choice: first, it reuses the infrastructure built for eager mode; second, XLA has been battle-tested and is sufficiently stable; third, Inductor’s success largely relies on Triton, which is difficult to port to the TPU architecture. This indirectly highlights how deeply PyTorch is coupled with GPUs, extending down to the compiler layer.

Calling torch.compile on TPU simply requires specifying backend="tpu".

PyTorch connects to the OpenXLA compilation stack via StableHLO, sharing the same lowering path as JAX and Pallas kernels.

Regarding distributed training, TorchTPU implements TPU collective communications at the Process Group level. DDP, FSDP2, and DTensor work out of the box. It also supports "asymmetric" patterns where rank 0 processes handle additional tasks like logging—something not possible with PyTorch/XLA, which only supports SPMD. According to official PyTorch/XLA documentation, TorchTPU will replace PyTorch/XLA upon its public release.

Using Streams to offload collective communications to SparseCore while performing matrix multiplications on TensorCore.

Performance Report Highlights

Coverage is the first metric. Google continuously tests approximately 5,000 of the most downloaded models from Hugging Face on TPUs. As of the report's writing, this covers 4,917 models across more than 500 architecture families. Of these, 98% run without crashing, and among 657 representative architectures, 97% produce outputs numerically consistent with GPU results.

Summary of Hugging Face model coverage.

The second metric is performance without specialized optimizations. Google wrote a 30B-parameter MuseGlimmer model using only standard PyTorch modules, without invoking any model-specific custom kernels. On a single TPU v7, prefilling with a sequence length of 5120 achieved 25,700 tokens per second, with an MFU (Model FLOPs Utilization) of 59%. Training on 4 TPU v7 chips yielded an MFU of 38%. Porting the equivalent implementation originally written for GPUs in TorchTitan with minor modifications resulted in a training MFU of 36%.

The third metric is scalability. For the Qwen3 235B MoE model, the PyTorch solution achieved approximately 94% of the throughput of Google’s own JAX-based MaxText solution on 64 TPU v7 chips. Scaling from 1,024 chips to 8,192 chips—a full TPU v7 Pod—resulted in only about a 6% drop in throughput per TensorCore, with MFU stabilizing between 27% and 29%. Loss curves remained nearly identical across all scales. The official blog further claims that training MFU for Qwen 3 8B can reach up to 52%.

Scaling data for Qwen3 235B MoE on 1,024 to 8,192 TPU v7 chips.

Total throughput scales nearly linearly with the number of chips.

Loss curves for the first 12 steps are nearly identical across different model scales.

In other words, the performance gap between PyTorch on TPU and JAX—Google's "native" framework—has been narrowed to single-digit percentages.

The inference-side data is equally impressive. According to SemiAnalysis InferenceX results cited in the report, TPU v7 paired with vLLM-TorchTPU achieves a cost of 2.01 cents per million tokens in an 8K input/1K output scenario, approximately 40% lower than the 3.6 cents reported during its initial debut in September.

The report also highlights several case studies: FLUX training time per step dropped from 11.49 seconds (out-of-the-box eager mode) to 2.78 seconds (with full compilation optimization); Meta’s DLRM v2 recommendation model achieved parity with Google’s highly optimized JAX baseline on 32 TPU v7x chips. According to the report, TorchTPU is already handling inference and training tasks in production environments for some major Google Cloud customers.

A "Full-Stack" Open Source Initiative

Beyond the technology itself, the completeness of the ecosystem carries significant weight.

• Inference: The vLLM backend developed in collaboration with Inferact and the SGLang backend with RadixArk retain upstream schedulers and API servers, requiring no changes to existing inference stacks. Additionally, Google, Red Hat, and IBM have released llm-d 1.0, enabling a unified inference entry point for both GPUs and TPUs on GKE.
• Training: TorchTitan provides pre-training recipes for dense models, MoE, and diffusion models; the reinforcement learning framework VeRL has integrated a TPU backend; and Ray supports scheduling different stages of the same workflow to CPUs, GPUs, or TPUs.
• Edge: LiteRT, formerly TensorFlow Lite, has officially joined the PyTorch ecosystem. Four core repositories—TorchTPU, vLLM-TorchTPU, SGLang-TorchTPU, and TorchTitan—have been open-sourced simultaneously.

Statements from industry partners are also noteworthy. Jana van Greunen, Engineering Director for PyTorch at Meta, stated frankly that Meta cannot be tied to a single hardware architecture. She noted that TorchTPU already supports critical production scenarios for their LLMs, advertising, and recommendation systems. It is worth remembering that Meta is a super-large customer of NVIDIA.

From Internal Codename to GitHub

Looking back at TorchTPU’s trajectory, the pace has been remarkably tight. The project launched internally at Google in mid-2025. In December of the same year, Reuters first disclosed the codename, reporting that Google was working closely with Meta and considering open-sourcing parts of the code. In April 2026, Google published the technical architecture on its developer blog, and the project lead presented at the PyTorch Conference Europe in Paris. By September, TorchTPU entered private beta and made its debut in SemiAnalysis’s inference benchmarks. Today, the entire stack has been officially open-sourced.

From "considering open-sourcing partial code" to fully publishing four core repositories and merging into upstream projects, Google’s answer has been more thorough than external expectations.

Will the CUDA Moat Be Filled?

A further question arises: Once the barrier at the framework layer is removed, what remains for NVIDIA?

First, it must be acknowledged that TorchTPU eliminates the framework-layer barrier, not the underlying hardware differences themselves. The report admits that simply adding .to("tpu") will not yield peak performance, as model shapes, sharding strategies, and operator choices may inherently favor GPUs. TPU TensorCores are essentially single-threaded processors requiring sufficiently wide shapes to saturate matrix units, which contrasts sharply with the GPU approach of running many small kernels in parallel.

Migrators should also note specific differences: TorchTPU defaults to BF16 multiplication with FP32 accumulation, whereas PyTorch on CUDA defaults to full FP32; default collective communication compilation introduces 30–40% end-to-end overhead, requiring additional annotations for optimization; each TPU can only be accessed by one process, meaning co-located training and inference in reinforcement learning workflows are not currently supported out-of-the-box. There remains a tuning effort required to bridge the gap between "running" and "running well."

Numerical precision comparison of PyTorch on CUDA vs. TPU

Second, this is not an isolated move by Google. At PyTorch Conference 2026, AWS also discussed the journey of bringing Trainium to native PyTorch. Its Neuron SDK has already discontinued support for training via PyTorch/XLA this year, shifting instead to the native TorchNeuron.

High-level architecture diagram of TorchNeuron

PyTorch is evolving into a de facto hardware abstraction layer, with PrivateUse1 becoming the standard slot for various chip vendors to integrate. When all hardware can run the same PyTorch code natively, the competitive focus shifts from "whether the ecosystem is usable" to "how much throughput you get per dollar"—precisely the metric Google repeatedly emphasizes in its announcement.

Of course, variables remain. Most of the performance data in the report was measured by Google itself and still requires independent validation by the community across more models and scenarios. Whether TPUs can reach ordinary developers with sufficient production capacity and pricing will also determine the ultimate impact of this open-source initiative. Google has stated that the priority for the next 6 to 12 months includes hardening for production environments, dynamic shape support, and FP8/FP4 quantization for next-generation TPUs.

But the direction is clear. For the past decade, choosing PyTorch essentially meant choosing CUDA. Starting today, that equation no longer holds.

References

https://developers.googleblog.com/torchtpu-running-pytorch-natively-on-tpus-at-google-scale/

https://www.denvergazette.com/2025/12/17/google-works-to-erode-nvidias-software-advantage-with-metas-help/

https://pytorchconferenceeu2026.sched.com/event/2Ikir

© THE END

Please contact this official account for authorization before reprinting.

For submissions or media inquiries: [email protected]