Google and Meta have joined forces to develop TorchTPU, an open-source initiative that aims to enable running PyTorch software without the need for intermediaries. The main goal of the initiative is to facilitate the development of tools that allow for transferring machine learning models from GPUs to TPUs faster than ever before, while retaining further benefits.
TorchTPU initiative, a partnership between Google and Meta, started its work in early 2025. The project has developed a new and very efficient backend that supports native PyTorch and allows for the implementation of various functions of existing models on TPUs with little effort. For instance, in order to transfer software from a GPU to a TPU, it would only be required to modify the device configuration, changing it from cuda to tpu. Furthermore, unlike other implementations, this project gives priority to the eager execution mode, which allows for performing operations without having to compile the whole script first.
TorchTPU Supports AI Model Training and Inference
The backend uses the XLA compiler stack from Google to turn the operations from PyTorch into the programmes that work on TPU hardware. Currently, XLA supports nearly 450 ATen operations, consisting of the operations you need for the forward and backward passes of the recent popular AI models. In addition, the compatibility with the libraries of Hugging Face has been tested, including Transformers, Diffusers, and Kernels. As a result, its library has almost 5000 downloads. Some popular AI models include Qwen, DeepSeek, Gemma, Llama, Mistral, Whisper, and BERT.
Besides, TorchTPU supports tools for distributed computing like DTensor, FSDP, vLLM, SGLang, TorchTitan, and Ray for tasks related to scaling and training. Developers may create custom TPU kernels with the help of JAX/Pallas and the experimental Helion programming language. It includes DeferAndFuse – the execution mode able to perform several operations at once, thus saving on processing. Google claims that this mode can provide 50% or more improvement in comparison to the standard eager execution in suitable cases. Currently, the system is under the process of development and is already in use by some major customers from Google Cloud.
The source code may be found on GitHub Repository.
