Google has open-sourced TorchTPU, a software stack designed to enable the training and execution of PyTorch models on Tensor Processing Units (TPUs). Developed in collaboration with the PyTorch team, the project aims to establish TPUs as a viable alternative to standard CPU and GPU hardware. Developers can migrate existing PyTorch applications with minimal code modifications while utilizing familiar interfaces for debugging, compilation, and distributed computing.
The stack supports step-by-step execution for debugging and torch.compile for performance optimization. Integration with the Hugging Face ecosystem is now available, specifically within Transformers 5, where users can specify 'device_map="tpu"'. Additionally, the Diffusers library now supports TPU-based execution and distributed computing. Google has also implemented automated testing to ensure ongoing compatibility between Transformers and TPU hardware.
For large language model inference, Google has introduced backends for vLLM and SGLang. SGLang utilizes KV cache reuse to optimize repeated requests, while the integration with llm-d allows for efficient request distribution across compute nodes based on existing context. These tools aim to streamline the deployment of LLMs on TPU infrastructure by managing input processing and response generation stages.
TorchTPU is integrated with TorchTitan to support distributed training of large neural networks and utilizes Pallas kernels. During testing with the Qwen 3 8B model, the system achieved up to 52% utilization of TPU capabilities and scaled across 8192 Ironwood TPUs. While these metrics reflect specific configurations, the open-source release is intended to lower the barrier for developers to adopt TPU hardware without requiring a total software stack overhaul.
