TorchTPU

// Overview

TorchTPU is an open-source PyTorch backend for Google Cloud TPUs, built by Google together with the PyTorch team at Meta. It adds a real `tpu` device to PyTorch, so most of the work of porting a GPU script is changing `device="cuda"` to `device="tpu"` and the distributed backend from `"nccl"` to `"tpu_dist"`. There is no JAX runtime underneath, and PyTorch's tensor view and in-place mutation semantics are preserved, with one process per device.

Execution is eager by default: operators are dispatched one at a time and run asynchronously. A debug mode (`TORCH_TPU_LAUNCH_BLOCKING=1`) makes them synchronous and adds runtime bounds checking, and a DeferAndFuse mode (`TORCH_TPU_DEFER_AND_FUSE=1`) fuses sequences of operators up to a materialization point such as `.cpu()` or `.item()` to cut HBM traffic. For compiled execution, a `"tpu"` Dynamo backend lowers FX graphs to StableHLO and hands them to the OpenXLA compiler, with a three-tier compilation cache spanning process memory, host shared memory and persistent remote storage.

TorchTPU is the successor to PyTorch/XLA. Google announced the project in April 2026 and said it would replace PyTorch/XLA once the public repository shipped; the repository and release 0.2.0 on PyPI arrived in October 2026. The feature-set page lists TPU v5e, v5p, v6e and v7 as validated, with v8i and v8t as near-term work, and marks checkpointing, activation offloading and multi-slice training up to 8,000 chips as validated. Custom kernels can be written with Pallas through `torch.tpu.pallas.jax_op`, with the experimental Helion DSL, or taken pre-built from Tokamax.

// What it does

• Native `tpu` device registered with PyTorch — no JAX runtime dependency, one process per device
• Three eager modes: strict (default, async), debug (synchronous with bounds checking), and DeferAndFuse (operator fusion)
• `torch.compile` support via a `"tpu"` Dynamo backend that lowers FX graphs to StableHLO for OpenXLA
• Distributed training at the `torch.distributed` process-group level: DDP, FSDP2, SimpleFSDP and DTensor
• Three-tier compilation cache across process memory, host shared memory and remote storage
• Custom kernels through Pallas, the experimental Helion DSL, or pre-built Tokamax kernels
• Example models in the repo: Llama, Qwen, ResNet, DLRM, minGPT and UNet

// Getting started

TorchTPU runs on a Cloud TPU VM. Install Python 3.12 and the wheel in that order so the PyTorch dispatcher and the TPU drivers line up, then switch your device string.

sudo apt update && sudo apt install -y python3.12 python3.12-venv
python3.12 -m venv env
source env/bin/activate

pip install torch_tpu

import torch

device = torch.device("tpu")  # Change from cuda to tpu

print(f"TPU Initialization complete. Primary device: {device}")

tpu_tensor = torch.ones((5, 5)).to(device)
direct_tpu_tensor = torch.randn((5, 5), device="tpu")

result_on_cpu = (tpu_tensor + direct_tpu_tensor).cpu()
print(f"Result Checksum (Sum): {result_on_cpu.sum():.4f}")

Commands and code are distilled from the project's own documentation — always check the official repo for the latest.

// When to use it

• Move an existing PyTorch training script from NVIDIA GPUs to Cloud TPUs with small edits
• Train large models on TPU pods with FSDP2, DTensor or multi-slice sharding
• Compile a PyTorch model through OpenXLA for TPU execution with `torch.compile`
• Write TPU kernels in Pallas or reuse Tokamax kernels from inside a PyTorch model

// How TorchTPU compares

TorchTPU alongside other open-source efficient training tools AI/TLDR tracks, ranked by GitHub stars.