Google and Meta Open Source TorchTPU as Native PyTorch Backend for TPUs

At the 2026 PyTorch North America Conference held in San Jose, Google and Meta open-sourced TorchTPU—a native PyTorch backend for Google's Tensor Processing Units (TPUs). In the October 21 keynote, Meta software engineer Zongwei Zhou stated that Meta (which created PyTorch in 2016) plans to use this backend to migrate larger-scale LLM workloads onto TPUs.

TorchTPU replaces PyTorch/XLA, a lazy tensor bridging library introduced in 2020. According to the announcement, that library "required invasive changes to source code." The newly launched backend aims to make model porting as simple as changing model.to("cuda") to model.to("tpu"). Google Cloud wrote in its announcement, "With TorchTPU, TPUs will join GPUs and CPUs as a native option for PyTorch." The announcement was co-authored by Bill Jia, Vice President of Core ML/AI Engineering.

After two talks laying out market demand, Jia and Zhou jointly delivered a keynote titled "Workload Interchangeability in the Agentic Era: Bringing Native PyTorch to Cloud TPUs." Mark Collier, Executive Director of the PyTorch Foundation, said that inference has overtaken training as the largest workload in data centers, giving rise to "enormous demand for tokens." Collier described the community's goal as developing open-source software that enables models to "run on any chip in any cloud and connect to any agent." The foundation's keynote slides illustrated the rapid pace of market development: DeepSeek's Flash model reduced costs by 60% in 37 days, while Kimi K3 improved runtime speed by 68% on the same hardware in under 33 days.

Mazin Gilbert, Executive Director of the Agentic AI Foundation, said that on February 6, 2026, agentic token consumption surpassed human consumption. Since then, agentic token consumption has grown 14x, while human consumption has grown only 2.8x. "So basically, this is a hockey-stick growth curve," Gilbert said. Before reusing this data, practitioners should verify its scope, as these public figures come from only one platform—OpenRouter.

The speakers said that because TorchTPU is integrated into PyTorch's low-level scheduler, Hugging Face models work out of the box. The vLLM and SGLang inference engines run natively, while TorchTitan supports dense, mixture-of-experts, and diffusion models for large-scale training. Ray and llm-d can scale TPU workloads across multiple hosts on Kubernetes. There are four open-sourced repositories: TorchTPU, vLLM-TorchTPU, SGLang-TorchTPU, and TorchTitan.

Jia demonstrated near-linear scaling with more than 8,000 TPU chips. As cluster size grows, per-chip throughput remains essentially constant. Google's published results include the Qwen3 235B mixture-of-experts model running on 8,192 TPU v7x chips with a model FLOPs utilization of 28%. The announcement candidly notes that this one-line code change is only a starting point:

This simple change alone will not reach peak TPU performance, because the model's structure, sharding, and operations may be better suited to GPU architectures; however, the model should work out of the box.

For this tuning gap, Google's solution is to add more agents. At the end of the talk, Jia announced the TPU Agent Suite. It is currently available in preview on GitHub and can be integrated into coding frameworks to automate environment onboarding, fault diagnosis, and performance tuning.

Zhou said Meta's PyTorch team developed TorchTPU jointly with Google "to make TPU a first-class native target for PyTorch," noting that this stems from production requirements for performance, reliability, and maintainability. "All of this is open source and available today, and this is just the beginning." In Google Cloud's announcement, Jana van Greunen, Director of Engineering for PyTorch at Meta, explained the motivation behind the move:

At Meta, advancing frontier AI requires enormous compute scale and rich hardware diversity. We cannot be locked into a single hardware architecture.

Google is not alone. AWS announced on October 20 the public testing of TorchNeuron, a native PyTorch backend for its Trainium chips, where code only needs to be changed to .to('neuron'). Both backends are registered through PrivateUse1, the extension point PyTorch provides for custom hardware. The first version of TorchNeuron still has some shortcomings: The Next Web reports that it does not yet support dynamic input shapes and skips about 22,000 PyTorch test cases. In an analysis piece written for Forkast, Blair Hayes points out that vendor lock-in has not disappeared, it has merely shifted:

Although PyTorch's extension points give developers a common language, the optimization layer and the supply layer remain firmly in the hands of vendors.

On Hacker News, the discussion thread about the TorchTPU release received 379 points. Commenters were both relieved and realistically skeptical. One commenter wrote, "I still remember the days of wrestling with XLA to get PyTorch running on TPUs, looks like we finally have official support!" Another asked where exactly one can buy a TPU. A commenter who had previously evaluated Google's on-premises TPU version on behalf of their company was more measured: "We evaluated these products but decided to stay with Nvidia, at least for another generation."

Portability is also a selling point outside the data center. In a session on deploying robot policies with ExecuTorch, Anthony Shoumikhin and Jacob Szwejbka from Meta noted that these policies can stay in PyTorch all the way from training to the device, without needing to convert tech stacks. The same quantization schemes and operator graphs can run on target platforms such as NVIDIA Jetson TensorRT and Intel OpenVINO.

The training track provided concrete figures from a cost perspective. Yueming Yuan from RadixArk, while introducing the Miles reinforcement learning framework, said that fusing two log-probability computations into a single kernel made this step more than 10 times faster. In one large-scale run, Delta updates reduced the amount of data that needed to be transferred per weight synchronization from 62 GB to under 1 GB.

Several sessions discussed open-source tools for training agents with reinforcement learning (RL). Anyscale's SkyRL separates training, inference, and environments, coordinating them through Ray. Hugging Face's OpenEnv (which joined the Agentic AI Foundation on October 22) provides a unified interface for each RL environment. Sergio Paniego Blanco from Hugging Face used TRL and OpenEnv to train a small Liquid AI model with four coding agents, raising its score on held-out tasks from 42% to 54%.

For teams considering a migration, changing just one line of device code is enough to get a model running (not yet tuned). So budget for sharding and kernel work, and confirm TPU capacity first. Developers who want to learn more can read the TorchTPU announcement, check out the TorchTPU codebase, or browse InfoQ's coverage of PyTorch Conference 2024. Session recordings will be posted on the PyTorch Foundation's YouTube channel within two weeks after the event.

Original link: https://www.infoq.com/news/2026/10/google-meta-torchtpu/