Google Opens TorchTPU for PyTorch Workloads on Its AI Chips

Google has released TorchTPU as an open-source path for running PyTorch workloads on its Tensor Processing Units (TPUs), giving AI developers another option for using cloud chips beyond Nvidia graphics processing units.

TorchTPU became available Oct. 20. It allows existing PyTorch workloads to run on Google’s TPUs with minimal code changes, with Google Cloud providing the computing infrastructure. Teams using PyTorch can test Google’s processors while continuing to work within the same framework.

Google’s TPU agent suite is in early preview. It is intended to help developers get started, diagnose problems and tune performance on TPUs. A separate integration, LiteRT-Torch, converts PyTorch models for use on devices such as phones and laptops.

Google worked with Meta and Hugging Face on compatibility with AI development tools. It also worked with Inferact on vLLM support and RadixArk on SGLang support, extending TorchTPU’s reach to tools used to run AI models.

Amazon announced its TorchNeuron backend for Trainium chips on Oct. 19.

Amazon has also contributed proposals and pull requests to the main PyTorch project, including work to remove hard-coded CUDA assumptions from tests. Those changes address specific references to Nvidia’s software in PyTorch’s development process.

PyTorch is an open-source framework for building, training and running AI models. Facebook developed it and released it in 2016, before changing its name to Meta. The company transferred the framework to the Linux Foundation in 2022.

On Oct. 20, Anyscale and Google Cloud also introduced native PyTorch support for Ray on Google TPUs. In that integration, Ray coordinates workers for training and model serving, while TorchTPU runs PyTorch operations on the chips.