rkj dev

Google Open-Sources TorchTPU for Native PyTorch Execution on TPUs

Co-developed with Meta, the open-source backend allows developers to run standard PyTorch models on Google TPUs with minimal code changes.

Illustration of a modern AI accelerator computing facility with connected server racks
Illustration: Modern AI data center clusters housing custom accelerator silicon for large-scale model training and inference.AI-generated illustration

In short

  • Google and Meta jointly launched TorchTPU, an open-source native PyTorch backend for Google TPUs published under the Apache 2.0 license.
  • TorchTPU replaces the older PyTorch/XLA lazy tensor approach with eager mode execution, torch.compile support, and roughly 450 ATen operations.
  • Integrations span Hugging Face Transformers v5, vLLM, SGLang, llm-d 1.0, Ray, and TorchTitan, where Google reported up to 52% MFU on 8,192 Ironwood TPUs.
  • The launch arrived one day after Amazon announced TorchNeuron for Trainium chips, intensifying cloud provider efforts to break Nvidia CUDA lock-in.

On October 20, 2026, at the PyTorch Conference, Google announced the open-source release of TorchTPU, a native PyTorch backend engineered to execute workloads directly on Google Cloud Tensor Processing Units. Co-developed with Meta's PyTorch engineering team since mid-2025, the software allows developers to port PyTorch models to TPUs without rewriting codebases, positioning Google's custom silicon alongside GPUs and CPUs as a first-class execution target.

The release directly tackles one of the AI sector's most persistent bottlenecks: the deep software lock-in anchoring developers to Nvidia hardware. According to a TechCrunch report cited by AI Weekly, approximately 85% of AI applications are built on PyTorch, which was originally created by Meta in 2016 and transferred to the Linux Foundation in 2022. Google's launch came exactly one day after Amazon announced its own open-source PyTorch initiative, TorchNeuron, aimed at porting models to AWS Trainium chips, as reported by NewsCord.

Moving Past PyTorch/XLA and Lazy Tensors

Historically, running PyTorch on Google TPUs required using PyTorch/XLA, a framework introduced in 2020 that relied on lazy evaluation. According to the PyTorch engineering team, PyTorch/XLA required intrusive code modifications to demarcate compilable code blocks, suffered from heavy compile-time latency, lacked an imperative eager execution mode, and struggled to integrate with modern distributed primitives.

Illustration of software engineers reviewing machine learning architectures in a tech workspace
Illustration: Engineers developing open-source machine learning runtimes and hardware execution layers.AI-generated illustration

TorchTPU abandons that legacy architecture. Instead of shoehorning PyTorch into JAX conventions, the project operates as a standard C++ dispatcher extension registered under PrivateUse1, featuring no JAX dependencies in the runtime. It supports eager execution out of the box through a newly introduced DeferAndFuse mechanism, alongside graph compilation powered by Torch Dynamo via torch.compile.

Under the hood, the backend implements approximately 450 PyTorch ATen operators, covering all Core ATen operations required for major deep learning architectures. It handles standard view semantics through decomposition and inversion while lowering collective communications directly to StableHLO rather than third-party communication stacks like NCCL. The software supports bounded dynamism, allowing tensor shapes to vary at runtime without forcing XLA recompilations as long as maximum bounds are defined.

Developer Workflow and Usability Changes

For machine learning practitioners, running existing models on TPUs now requires minimal refactoring. According to Google Cloud, developers can switch hardware targets by replacing device calls such as model.to("cuda") with model.to("tpu") or using PyTorch's accelerator autodetect APIs.

In the Hugging Face ecosystem, native support is baked directly into Transformers v5, allowing developers to target TPUs simply by setting device_map="tpu". The Diffusers library also supports eager mode, torch.compile, and tensor parallelism on TPUs. Google confirmed that 5,000 of the most downloaded Hugging Face models—spanning LLMs like Qwen, Llama, and Mistral, as well as vision and audio models such as Whisper, CLIP, and ViT—have been validated for out-of-the-box execution.

Hardware precision defaults have also been tailored to the TPU architecture. While CUDA defaults to standard FP32 operations, TorchTPU defaults to a medium precision mode where inputs convert to 16-bit bfloat16 for TPU systolic Matrix Multiply Units (MXUs) while accumulating in FP32. High and highest precision modes use multi-pass algorithmic decompositions (requiring 3× and 6× MXU passes, respectively), and Google warns they incur steep performance overheads.

High-Throughput Serving and Training at Scale

TorchTPU arrives paired with production integrations across prominent inference and training stacks. For inference serving, Google partnered with Inferact to build a native PyTorch backend for vLLM on TorchTPU, incorporating custom Pallas kernels, continuous batching, and multi-chip parallelism. A collaboration with RadixArk brings TPU support to SGLang, utilizing RadixAttention for prompt KV cache reuse across multi-step agent workloads.

Illustration of high-throughput distributed model inference across server clusters
Illustration: High-throughput orchestration routing AI inference queries across distributed server nodes.AI-generated illustration

To orchestrate distributed inference, Google released llm-d 1.0 in partnership with Red Hat, IBM, and the Cloud Native Computing Foundation (CNCF). According to Google, llm-d acts as a prefix-cache-aware Kubernetes router separating prefill and decode stages. Zenteiq, an early adopter using llm-d on Google Kubernetes Engine (GKE), reported achieving over 4× higher aggregate inference throughput, boosting daily synthetic data generation from 1.47 billion to 5.97 billion tokens per 1,000 TPU v6e accelerators.

On the distributed training side, TorchTPU integrates with TorchTitan, Meta's reference library for large-scale training. Google reported that TorchTitan running on TorchTPU achieved up to 52% Model Flops Utilization (MFU) on Qwen 3 8B, demonstrating stable near-linear scaling across clusters of up to 8,192 Ironwood TPUs (TPU v7x). TorchTPU also upstreamed support to the VeRL reinforcement learning framework and added native scheduling capabilities to Ray for heterogeneous compute pipelines.

Beyond cloud data centers, Google announced that LiteRT (formerly TensorFlow Lite) has joined the PyTorch ecosystem. Developers can convert PyTorch models using a single litert-torch.convert() function to deploy onto Android, iOS, Windows, and embedded hardware.

Hardware Diversity in the Enterprise

The project is available on GitHub under the Apache 2.0 license, requiring Python 3.11 or newer. Pre-built wheels are published via PyPI, and developers can build from source using Bazel.

Meta's director of PyTorch engineering, Jana van Greunen, noted that the collaboration allows Meta to deploy PyTorch workloads—including large language models, ad recommendation systems, and core production pipelines—directly to Google TPUs, ensuring the company is not constrained by a single hardware architecture.

While Google and Amazon are actively lowering switching costs by decoupling PyTorch from Nvidia CUDA primitives, industry analysts emphasize that software availability alone does not guarantee widespread migration. As noted by aipolix.com, Google's scaling and MFU figures have not yet been independently verified outside commercial partnerships, and real-world adoption will hinge on whether running production workloads on TPUs delivers superior cost-per-token economics compared to established GPU infrastructure.

Questions people ask

What is TorchTPU?

TorchTPU is an open-source PyTorch backend developed jointly by Google and Meta that allows standard PyTorch code to run natively on Google Cloud TPUs.

How does TorchTPU differ from PyTorch/XLA?

Unlike PyTorch/XLA, which relied on lazy tensors and required manual code annotations, TorchTPU supports native eager execution, torch.compile via Torch Dynamo, and standard ATen operations without JAX runtime dependencies.

Which libraries and frameworks support TorchTPU at launch?

TorchTPU integrates out of the box with Hugging Face Transformers v5 and Diffusers, vLLM, SGLang, llm-d 1.0, TorchTitan, Ray, and the VeRL reinforcement learning framework.

What license and Python versions does TorchTPU require?

TorchTPU is licensed under the Apache 2.0 open-source license and requires Python 3.11 or newer.

Sources

  1. TorchTPU is now open source: Run native PyTorch on TPUsGoogle Cloud, Oct 20, 2026, official source
  2. google-pytorch/torch_tpuGitHub, Mar 9, 2026, official source
  3. Introducing TorchTPUPyTorch, Oct 20, 2026
  4. Amazon Launches TorchNeuron to Challenge Nvidia’s AI Chip DominanceNewsCord, Oct 21, 2026
  5. Google open-sources TorchTPU to bring native PyTorch support to its TPUsCrypto Briefing, Oct 20, 2026
  6. Google opens TorchTPU to pull PyTorch developers off CUDAAI Weekly, Oct 20, 2026
  7. Google open-sources TorchTPU, making its TPUs a native PyTorch device built with Meta, Hugging Face, vLLM and SGLang supportaipolix.com, Oct 20, 2026

How this story was made: the newsroom picked it up from Google Search, Hacker News and Techmeme, gathered the full text of the sources above, and drafted it with AI assistance. Every factual claim was then checked against those sources before publishing (41 claims checked). Illustrations marked as AI-generated are not photographs. Spotted an error? Tell us.

#TorchTPU #Google Cloud #PyTorch #Meta #AI Chips #Machine Learning

Published October 21, 2026 at 01:05 UTC