Google Open-Sources TorchTPU: A Native PyTorch Backend for TPUs

Google has released TorchTPU, a specialized backend designed to integrate PyTorch seamlessly with Google's Tensor Processing Units (TPUs). This project allows developers to execute PyTorch operations and models directly on TPU hardware, leveraging high-performance acceleration.

🚀 Core Capabilities

TorchTPU isn't just a wrapper; it provides deep integration through several key components:

• Custom ATen Kernels: Specialized implementations for various PyTorch operations to maximize TPU efficiency.
• XLA Compilation Cache: An optimization layer that caches computations to speed up execution.
• Distributed Scaling: Native support for advanced training paradigms, including:
  1. Data Parallelism
  2. Tensor Parallelism

Supported Model Examples

The repository provides reference implementations for a variety of architectures:

---

📖 Documentation & Learning

Users can access the comprehensive User Guide. For those preferring a hands-on approach, interactive notebooks are available.

[!TIP] To run these notebooks, use marimo edit file.py within your browser. For the best experience, deploy these on a TPU VM (a v6e GCE instance is the recommended choice).

---

🛠️ Installation Guide

Environment Setup

TorchTPU follows the PyTorch Release Compatibility Matrix regarding Python version support.

Step 1: Initialize a Virtual Environment You can use standard venv or the faster uv tool:

# Using standard venv
python -m venv .venv
source .venv/bin/activate

# OR using uv
uv venv .venv
source .venv/bin/activate

Step 2: Install via Pip To install the current public release (version 0.2.0), use the following command. Note that we specify the CPU index to avoid default CUDA installations.

pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpu

---

🏗️ Building from Source

For developers needing to modify the backend, TorchTPU utilizes Bazel.

The Build Process

The standard command to generate a wheel is: bazel build --config=wheel_common //ci/wheel:torch_tpu_wheel

Optimization Levels

• Optimized (-c opt): The default for all OSS Bazel builds.
• Unoptimized (-c fastbuild or -c dbg): Used for debugging C++ code.

[!WARNING] Unoptimized builds will significantly degrade performance. To build an unoptimized wheel, you must explicitly pass the flag --//:allow_unoptimized_wheel=True. Failure to do so will result in a build-time error.

Technical Nuance: The wheel_common Config

The wheel_common configuration is critical because it sets --//:wheel_build=True. This ensures the shared XLA/MLIR backend is routed through a torch_version -reset transition.

$\text{Result} \rightarrow \text{Single } \texttt{libxla\_base.so}$

Without this, each version would ship its own backend copy, causing the import torch command to crash due to duplicate static registrations.

Remote Build Execution (RBE)

By default, wheel_common enables the RBE remote cache and executor. If you lack RBE credentials, disable it using the following sequence:

bazel build --config=wheel_common //ci/wheel:torch_tpu_wheel --config=no_rbe

Note: The order is vital; --config=no_rbe must come last to override the remote-cache flags.

Finalizing Source Installation

Once the wheel is built, install it while ensuring the CPU version of PyTorch is used:

pip install ../bazel-bin/ci/wheel/*.whl --extra-index-url https://download.pytorch.org/whl/cpu

---

⚙️ Dependency Management & Tooling

Bazel & Bazelisk

The project relies on Bazelisk, a wrapper that automatically manages Bazel versions.

Installation Options:

• Via Go: go install github.com/bazelbuild/bazelisk@latest
• Via Binary: wget -O ~/bin/bazel https://github.com/bazelbuild/bazelisk/releases/download/v1.27.0/bazelisk-linux-amd64

After installation, ensure the binary is in your PATH and verify with: bazel info

IDE Integration

For those navigating the C++ source code, VSCode with the clangd extension is highly recommended.

---

🗺️ Workflow Summary

✅ Getting Started Checklist

• Review Official PyTorch Documentation
• Explore TorchTPU Tutorials
• Study the Public APIs
• Verify installation with bazel info