DEV Community

Cover image for Google and Meta Open-Source TorchTPU for Running PyTorch on TPUs
The daily flare
The daily flare

Posted on Originally published at thedailyflare.com

Google and Meta Open-Source TorchTPU for Running PyTorch on TPUs

Originally published on The Daily Flare.

Key takeaways

  • Google and Meta released TorchTPU as an open-source PyTorch backend for Google TPUs.
  • The stack connects with tools including TorchTitan, vLLM, SGLang and Hugging Face libraries.
  • A model may run with a small device-setting change, but peak performance can require further tuning.

Google and Meta have released TorchTPU, an open-source backend that lets developers run PyTorch models on Google's TPU accelerators with limited code changes.

Announced at PyTorch Conference North America in San Jose, TorchTPU gives TPUs a native place alongside CPUs and GPUs in PyTorch workflows. The project is intended to make it easier to move existing models to Google's chips without rewriting them for another machine-learning framework. Google Cloud's official announcement says developers can use familiar PyTorch APIs for training, serving and deployment.

A native backend for PyTorch workloads

TorchTPU was co-developed by Google and Meta's PyTorch team. InfoQ reported that Meta plans to use the backend for larger-scale language-model workloads on TPUs. Meta's Jana van Greunen, director of PyTorch engineering, said the collaboration supports the company's need to work across different hardware architectures and to run production systems for language models, advertising and recommendations.

The backend is designed to work with PyTorch's eager execution mode for debugging and with torch.compile for optimization. According to Google, it also supports standard APIs and distributed computing tools. Because it connects through PyTorch's low-level dispatcher, the conference speakers said, models from Hugging Face can work without extensive changes. Google also said its Transformers library supports sending models to a TPU with device_map="tpu", while its Diffusers library can run generative models on the hardware.

TorchTPU is presented as a successor to the older PyTorch/XLA approach, which InfoQ described as a lazy-tensor bridge that could require intrusive changes to source code. With the new backend, the basic migration may involve changing a model's device setting from CUDA to TPU. That does not guarantee peak speed: Google and PyTorch have cautioned that model shapes, sharding strategies and operations may need adjustment to suit TPU hardware.

Training and serving tools join the TPU stack

The release includes integrations for several parts of the model-development pipeline. Google and Meta listed TorchTitan for large-scale training, and vLLM and SGLang for inference. Ray and llm-d can help coordinate TPU workloads across multiple machines in Kubernetes environments. InfoQ reported that four code repositories were opened: TorchTPU, vLLM-TorchTPU, SGLang-TorchTPU and TorchTitan.

Google's conference presentation included a Qwen3 235B mixture-of-experts model running across 8,192 TPU v7x chips, with reported model FLOPs utilization of 28%. The presentation also showed throughput per chip staying roughly steady as the cluster expanded beyond 8,000 chips, which Google characterized as near-linear scaling. Those figures describe a specific demonstration, not a guarantee that every model will achieve similar results.

For training, TorchTitan provides recipes for dense models, mixture-of-experts systems and diffusion models. Google's developer documentation describes TPU-specific options for attention, projection and loss operations, including custom kernels intended to reduce memory use. These are examples of the additional engineering that can be needed after a model first runs on the device.

Google also announced a preview of a TPU agent suite for coding tools. The company says the agents are intended to help with onboarding, diagnostics and performance tuning. They address a practical gap between making a model run and adapting it to use the hardware efficiently.

Hardware choice grows, while tuning remains essential

TorchTPU arrives as other cloud providers add native PyTorch paths for their own processors. The PyTorch Foundation's conference recap also highlighted TorchNeuron, an open-source backend for AWS Trainium that entered public beta during the conference. Both projects use PyTorch's extension mechanisms to connect the framework to specialized chips, while leaving hardware-specific optimization in each vendor's stack.

The wider context is a growing demand for computing capacity to serve models. At the conference, PyTorch Foundation executive director Mark Collier said inference had become the largest data-center workload, while Agentic AI Foundation executive director Mazin Gilbert cited rising token use by AI agents. InfoQ noted that some of the figures discussed came from OpenRouter, a single platform, so they should not be treated as a complete measure of industry-wide activity.

For developers, TorchTPU's main promise is a more familiar route to experimenting with Google's accelerators using PyTorch tools. The initial device change may ease migration, but production performance still depends on model-specific work, distributed setup and tuning. The project makes the software entry point more direct; it does not remove the differences between hardware platforms.

Top comments (0)