Google has open-sourced TorchTPU, a software layer that lets developers run PyTorch—the widely used toolkit for building and running AI models—on Google’s Tensor Processing Units, or TPUs. Developed with Meta’s PyTorch team, it is designed to move existing workloads to TPU hardware with relatively small code changes rather than a rewrite.
Google says TorchTPU supports both eager execution, which runs operations as they are called and is useful for debugging, and compiled execution for production workloads. Developers can switch a device setting from “cuda” to “tpu” and use the TPU distributed backend; Google also points to PyTorch’s accelerator-detection API as an alternative. TorchTPU is available on GitHub, but its roadmap still includes work on dynamic input shapes, compilation speed, operator coverage, reinforcement learning and quantization.
What changes for developers?
The launch connects TPUs to software many PyTorch developers already use. Google says TorchTPU works with serving tools vLLM and SGLang, and training and orchestration projects including TorchTitan, Ray and TorchRec. In Transformers v5, Hugging Face says developers can select a TPU with the setting device_map="tpu". Its Diffusers library also supports TPU execution, including compiled mode and splitting work across multiple accelerators.
This is a different approach from PyTorch/XLA, Google’s earlier route to running PyTorch on TPUs. Google says that system required more intrusive code changes, had long compilation times and did not support PyTorch’s eager mode. Google has built a test pipeline that, according to InfoQ, evaluates about 5,000 Hugging Face models across more than 500 architectures. That signals broad testing, but does not mean every model or workload is guaranteed to work unchanged.
How does it compare with other hardware options?
Google reports that TorchTPU’s DeferAndFuse mode—a hybrid of eager and compiled execution—can be at least 50% faster than its Strict Eager mode while reducing peak accelerator memory use. In a separate test, Google ran a 30-billion-parameter Muse Glimmer model in standard PyTorch: InfoQ reports inference throughput of 25,700 tokens per second on one TPUv7 chip at a sequence length of 5,120, and 22,400 at 10,240. Training on four TPUv7 chips reached 5,600 and 4,600 tokens per second at those respective lengths. These are Google-reported results, not an independent comparison.
TechCrunch places TorchTPU alongside Amazon’s TorchNeuron, announced the same week for running PyTorch workloads on Amazon’s Trainium chips. Both aim to make it easier to use alternatives to Nvidia GPUs. TorchTPU’s release adds another hardware option for PyTorch users, but does not itself show that moving a workload will lower costs or outperform a developer’s current setup.
Why it matters
For developers already using PyTorch, TorchTPU could make it easier to test or deploy workloads on Google TPUs without rebuilding them from scratch. Whether that switch improves cost or speed for a particular model remains unclear; the performance figures cited are Google’s own results.
