Google open-sourced TorchTPU, so PyTorch on TPUs just got a lot easier
Google released TorchTPU as open source, giving PyTorch developers a native way to run models on the company’s Tensor Processing Units. Google also announced TPU support in the vLLM and SGLang serving engines, the TorchTitan training framework, Ray, and Hugging Face’s libraries.
Teams that write their models in PyTorch and run them on GPUs by default have the most to gain. Google says those workloads now move to TPUs, its custom AI chips, with minimal code changes. If that claim holds up in practice, a team weighing GPU costs can test Google’s hardware without first rebuilding its software.
What developers change
TorchTPU follows the standard PyTorch workflow. Developers debug in eager mode, where each operation runs as soon as it is called. They then switch to torch.compile, which turns the model into an optimized graph for production. Distributed training primitives come built in.
Hugging Face cut the switch down to one argument. In Transformers v5, passing device_map=”tpu” sends a model to a TPU, and Google runs the Transformers test suite on TPUs every night. The Diffusers library for generative models also runs natively, with tensor parallelism across multiple chips.
Meta says it already runs LLM, ads and recommendation workloads on TPUs through TorchTPU in key production use cases. “We cannot be constrained by a single hardware architecture,” said Jana van Greunen, Meta’s director of PyTorch engineering.
Google’s numbers
TorchTitan training on TorchTPU reached up to 52% model FLOPs utilization on Qwen 3 8B, with near-linear scaling across as many as 8,192 Ironwood TPUs. Near-linear scaling means throughput grew almost in step with chip count. Model FLOPs utilization measures how much of a chip’s theoretical math capacity goes into useful training work, so 52% means about half.
A separate result comes from llm-d, an open-source router that spreads inference requests across clusters. Google released llm-d 1.0 with Red Hat and IBM. Zenteiq.ai says llm-d on Google Kubernetes Engine raised its daily output from 1.47 billion to 5.97 billion tokens per 1,000 TPU v6e chips, roughly a fourfold gain. Zenteiq credits the gain to routing that sends each request to a replica already holding its context.
Still in preview
Two pieces remain in preview. One is the TPU backend Google contributed to VeRL, a widely used reinforcement-learning framework. The other is the TPU agent suite, which plugs into AI coding tools to automate onboarding, diagnostics and inference tuning for PyTorch models. Google also brought LiteRT, formerly TensorFlow Lite, into the PyTorch ecosystem for running models on phones and other devices.
The TorchTPU code is available for free on GitHub.
Must read:
- 20 open-source cybersecurity tools to keep your team ready for anything
- GitHub CISO on security strategy and collaborating with the open-source community
Subscribe to the Help Net Security ad-free monthly newsletter to stay informed on the essential open-source cybersecurity tools. Subscribe here!
