Google has released TorchTPU, an open-source PyTorch backend that allows users to run PyTorch models and operations directly on Google’s Tensor Processing Units (TPUs). The project, available on GitHub, provides PyTorch developers with native TPU support without requiring external compilation frameworks.
According to the repository, TorchTPU includes PyTorch ATen kernel implementations for TPU, enabling efficient execution of PyTorch operations on TPU hardware. The library incorporates a compilation cache for optimizing XLA computations, which are the underlying mathematical operations that power machine learning models.
Key features include support for distributed training approaches such as Data Parallelism and Tensor Parallelism, allowing users to scale model training across multiple TPUs. The backend integrates with torch.compile() for graph mode execution, a PyTorch feature that optimizes model performance by converting Python code into optimized computation graphs.
TorchTPU includes examples for training and inference across various model architectures, including Llama, Qwen, ResNet, DLRM, minGPT, and UNet. The library also provides debugging and benchmarking utilities to help developers optimize their code.
The current public release is version 0.2.0, available on PyPI. Installation requires Python 3.11 or later, matching PyTorch’s official compatibility matrix. Users can install via pip with the command: pip install torch_tpu==0.2.0 –extra-index-url https://download.pytorch.org/whl/cpu.
Developers looking to build from source can use Bazel, the dependency management system employed by the project. The repository includes comprehensive setup instructions and documentation at https://torchtpu.readthedocs.io/, featuring interactive Marimo notebooks (similar to Jupyter) for learning the platform. Google notes that running these notebooks requires access to a TPU VM or v6e GCE instance.
The project is licensed under Apache 2.0 and welcomes external contributions. Developers interested in contributing must sign Google’s Contributor License Agreement and follow the guidelines outlined in the repository.
Key facts
- TorchTPU enables native PyTorch execution on Google’s Tensor Processing Units
- The library includes PyTorch ATen kernels and a compilation cache for optimized XLA computation
- Supports distributed training through Data Parallelism and Tensor Parallelism
- Integrates with torch.compile() for graph mode execution
- Version 0.2.0 is available on PyPI; requires Python 3.11 or later
- Includes example implementations for Llama, Qwen, ResNet, DLRM, minGPT, and UNet models
- Licensed under Apache 2.0 with documentation and interactive notebooks available