Google open-sources TorchTPU, a native PyTorch backend for TPUs

Google has published the source code of torch_tpu, a PyTorch backend that makes it possible to run the framework's models and operations directly on its Tensor Processing Units (TPUs). Until now, pairing PyTorch with Google's in-house accelerators has meant taking roundabout routes; the project offers a more direct path.

What's inside

The project ships a custom implementation of ATen kernels for a range of PyTorch operations, a compilation cache that speeds up execution of XLA computations, and integration with torch.compile() for graph-mode execution. Distributed training is supported, covering both data parallelism and tensor parallelism.

The repository bundles examples for several popular models: Llama, Qwen, ResNet, DLRM, minGPT and UNet. There are also utilities for debugging and benchmarking.

Installation and building

The current public release on PyPI is 0.2.0; Python 3.11 and above is supported, with a version policy that mirrors PyTorch's own compatibility matrix. The package installs via ordinary pip, though the docs strongly advise adding an extra index for the CPU build of torch — otherwise the CUDA variant is pulled in by default.

Building from source relies on Bazel (via the Bazelisk wrapper), and here lie a few subtleties. The wheel build requires the wheel_common config, without which it simply fails. The reason is technical: without that flag, every version carries its own copy of the XLA/MLIR backend, and the resulting wheel then aborts right on import torch because of duplicate static registrations. The documentation even warns about flag ordering — --config=no_rbe must come after --config=wheel_common, since the config listed later is the one that wins.

Documentation

A user guide is hosted on readthedocs, and the interactive notebooks are built with Marimo (a Jupyter-like tool) and live in the docs folder. The easiest way to run them is directly on a TPU instance in Google's cloud.

The project is licensed under Apache 2.0. At the time of writing, the Hacker News discussion had gathered 189 points and 47 comments.

Source: Hacker News