Google Opens TorchTPU for PyTorch on TPUs

Google open-sourced TorchTPU, a PyTorch software stack for TPUs, at PyTorch Conference 2026, held on the 20th local time. Co-developed with Meta's PyTorch team, the tool keeps the existing PyTorch API and supports model training, inference and deployment on TPUs. Developers can run existing applications on TPUs with minimal code changes.

TorchTPU provides a standard execution mode for checking model errors, performance optimization with “torch.compile,” distributed training and diagnostic tools. Google worked with Hugging Face to enable models to be placed on TPUs in Transformers v5 using the “device_map=\"tpu\"” setting. The image-generation tool Diffusers supports standard execution, compiler-based optimization and tensor parallelism, which splits a model across multiple chips. Google also set up ongoing tests of Hugging Face models on TPU infrastructure.

Google also built TPU backends for the large-model inference tools vLLM and SGLang in collaboration with Inferact and RadixArk. These backends optimize computation for TPUs while retaining existing schedulers and API servers. The Kubernetes-based system llm-d, which distributes inference requests across multiple servers, was released as version 1.0. It routes requests to servers with cached context and separates prefill, the processing of input, from decoding, the generation of answers. On Google Kubernetes Engine (GKE), GPUs and TPUs can be used through a common interface.

TPU support has also been integrated into the training framework TorchTitan. Google said model FLOPs utilization (MFU), the share of theoretical computing performance used for actual training, reached up to 52% when training Qwen 3 8B. It also said performance rose almost in proportion to the number of chips when scaling to as many as 8192 Ironwood TPUs. The reinforcement-learning framework VeRL gained a TPU backend and reference configuration. Ray supports assigning CPUs, GPUs and TPUs to data processing, training, reinforcement learning and inference within a single workflow.

Google made a preview of the TPU Agent Suite available on GitHub to help with initial setup, error diagnosis, usage questions and inference performance tuning. It also released LiteRT-Torch, which connects PyTorch models to the on-device inference tool LiteRT. The tools convert models to run on Android, iOS, Windows, macOS, Linux, the web and Internet of Things devices, using each device's CPU, GPU or neural processing unit (NPU).

Amazon announced TorchNeuron the previous day, a tool for running PyTorch code on its Trainium AI chips, and said it would be available in public beta from the 20th. Amazon said users can switch to Trainium by changing one setting, without rewriting code. PyTorch is a widely used open-source framework for developing and training AI models. Because much code and tooling is tailored to Nvidia GPUs and CUDA, moving to other chips has required additional work. Google and Amazon are seeking to reduce that migration burden and broaden use of their own chips.

Google open-sourced TorchTPU, a PyTorch software stack for TPUs, at PyTorch Conference 2026, held on the 20th local time. Co-developed with Meta's PyTorch team, the tool keeps the existing PyTorch API and supports model training, inference and deployment on TPUs. Developers can run existing applications on TPUs with minimal code changes.

TorchTPU provides a standard execution mode for checking model errors, performance optimization with “torch.compile,” distributed training and diagnostic tools. Google worked with Hugging Face to enable models to be placed on TPUs in Transformers v5 using the “device_map=\"tpu\"” setting. The image-generation tool Diffusers supports standard execution, compiler-based optimization and tensor parallelism, which splits a model across multiple chips. Google also set up ongoing tests of Hugging Face models on TPU infrastructure.

Google also built TPU backends for the large-model inference tools vLLM and SGLang in collaboration with Inferact and RadixArk. These backends optimize computation for TPUs while retaining existing schedulers and API servers. The Kubernetes-based system llm-d, which distributes inference requests across multiple servers, was released as version 1.0. It routes requests to servers with cached context and separates prefill, the processing of input, from decoding, the generation of answers. On Google Kubernetes Engine (GKE), GPUs and TPUs can be used through a common interface.

TPU support has also been integrated into the training framework TorchTitan. Google said model FLOPs utilization (MFU), the share of theoretical computing performance used for actual training, reached up to 52% when training Qwen 3 8B. It also said performance rose almost in proportion to the number of chips when scaling to as many as 8192 Ironwood TPUs. The reinforcement-learning framework VeRL gained a TPU backend and reference configuration. Ray supports assigning CPUs, GPUs and TPUs to data processing, training, reinforcement learning and inference within a single workflow.

Google made a preview of the TPU Agent Suite available on GitHub to help with initial setup, error diagnosis, usage questions and inference performance tuning. It also released LiteRT-Torch, which connects PyTorch models to the on-device inference tool LiteRT. The tools convert models to run on Android, iOS, Windows, macOS, Linux, the web and Internet of Things devices, using each device's CPU, GPU or neural processing unit (NPU).

Amazon announced TorchNeuron the previous day, a tool for running PyTorch code on its Trainium AI chips, and said it would be available in public beta from the 20th. Amazon said users can switch to Trainium by changing one setting, without rewriting code. PyTorch is a widely used open-source framework for developing and training AI models. Because much code and tooling is tailored to Nvidia GPUs and CUDA, moving to other chips has required additional work. Google and Amazon are seeking to reduce that migration burden and broaden use of their own chips.