Google released TorchTPU, its version of the PyTorch framework, as an open-source project on Tuesday. It lets developers move existing AI applications to Google’s chips without rewriting their code. Alongside it, the company introduced a migration assistance tool and updated LiteRT’s support for PyTorch applications.
Moving to TPUs with a configuration change
TorchTPU runs AI applications on Tensor Processing Units (TPUs), Google’s own AI chips. Moving existing PyTorch applications to TPUs, and consequently to Google Cloud, requires a simple configuration change.
PyTorch is an open-source toolkit for building, training and running AI models. Its ties to Nvidia hardware have historical roots: the company’s GPUs were already standard equipment in AI research when the framework was being developed, so many features were built specifically for them.
Development tool compatibility and migration assistance
Google says it worked directly with Meta, Hugging Face and other partners on compatibility for TorchTPU’s features and libraries. Those partners included Inferact, the developer of vLLM, and RadixArk, which makes SGLang. The collaboration aimed to make TorchTPU work with these development tools much as the original PyTorch does.
The newly launched TPU agent suite is designed to help developers move applications to Google’s chips and cloud. The dedicated tool is also intended to help fix errors encountered during migration.
The announcement also includes an update to LiteRT, a tool for running AI on phones, laptops and embedded devices. The update improves its compatibility with AI applications built in PyTorch.
Amazon opened a path to Trainium a day earlier
Google’s TorchTPU release follows a similar move by Amazon a day earlier. On Monday, Amazon released its own version of PyTorch, called TorchNeuron, which lets developers move AI applications to Amazon Trainium chips and AWS without rewriting them.
Amazon also said it is submitting code to the main PyTorch project that separates various features from Nvidia’s chips. The change is intended to make the framework more open to other hardware by default.
