TorchTPU PyTorch TPU Guide: Install and Migrate in 2026

TorchTPU makes a cuda-to-tpu PyTorch port relatively small, but TPU access, package selection, compilation, and tensor shapes determine whether it works well.

TorchTPU’s actual promise

TorchTPU is a native PyTorch backend for Google TPUs, not a hosted model API or a replacement for a TPU VM. The PyTorch announcement describes eager execution, torch.compile, distributed primitives, and support for roughly 450 ATen operations; the public repository lists ATen kernels, compilation caching, distributed examples, and torch.compile integration.

The migration target is deliberately familiar:

# Before
model.to("cuda")

# First TPU port
model.to("tpu")

That is a portability test, not a performance guarantee. PyTorch’s announcement notes that peak TPU performance still depends on model shapes, sharding, and operations suited to the hardware.

What you need before installing

TorchTPU is useful only where a Google TPU is provisioned. The repository recommends running its notebooks on a TPU VM and calls a v6e Google Compute Engine instance the easiest starting environment. A CPU-only laptop can install packages, but it cannot verify TPU execution.

The public README supports Python 3.11+ and identifies torch_tpu==0.2.0. Use an isolated environment so the experiment does not inherit a CUDA build or unrelated dependencies.

Install torch_tpu without pulling CUDA

Create a virtual environment, then install the CPU PyTorch wheel with the pinned TorchTPU package:

python -m venv .venv
source .venv/bin/activate

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

Per the README, omitting the extra index can pull a CUDA PyTorch wheel, the wrong dependency for a TPU VM.

For a first migration, use the released wheel. Source builds require Bazel or Bazelisk and the repository’s documented --config=wheel_common and, where needed, --config=no_rbe flags.

Verify the backend with a minimal PyTorch program

Run this on the TPU VM:

import torch

x = torch.randn(2, 2, device="tpu")
y = torch.randn(2, 2, device="tpu")
z = x @ y

print(z)
print(z.device)

A successful result with a tpu device confirms tensor allocation and matrix multiplication reached the TPU backend. If the device is unavailable, check the VM runtime and package installation before changing model code.

TorchTPU supports eager execution for debugging and torch.compile for compiled execution. The PyTorch announcement reports that its DeferAndFuse mode can be 50% or more faster than strict eager mode in project testing; treat that as a project claim, not a universal benchmark.

Port a real PyTorch workload

Start with the smallest working change, then add TPU-specific distributed configuration only after the single-device path is verified.

1. Replace the device selection with torch.device("tpu") or use torch.accelerator.current_accelerator() where supported.
2. Move the model and input tensors to the same device.
3. Run one short batch before enabling compilation or distributed execution.
4. Keep batch sizes and tensor shapes stable while measuring; shape changes can trigger compilation work.
5. Test the representative model path and distributed pattern before committing to a larger migration.

A distributed CUDA snippet commonly uses NCCL. TorchTPU’s published example uses the TPU process-group backend instead:

import torch
import torch.distributed as dist

device = torch.device("tpu")
dist.init_process_group(backend="tpu_dist")
model = model.to(device)

The exact distributed launch command depends on the TPU environment and workload. Changing the device string does not also configure sharding, process placement, or input pipelines.

Expect the first run to differ from later runs. TorchTPU’s documented compilation cache can reuse repeated graphs, so record warm-up time separately from steady-state throughput.

Where TorchTPU is ready—and where it is still rough

Google documents TorchTPU integrations with vLLM, SGLang, TorchTitan, and Hugging Face tooling. Its inference announcement describes native TPU paths for vLLM and SGLang, including KV-cache movement and serving features. Google’s Cloud announcement reports the 52% MFU Qwen 3 8B TorchTitan result and production use claims.

The practical limits are less polished than the announcement headlines:

The decision rule for adopting TorchTPU

TorchTPU is worth testing now if you already have Google TPU access, run PyTorch models, and want a common framework path across training or inference. The strongest early use case is a team willing to profile and tune shapes rather than expecting a one-line port to deliver GPU-equivalent efficiency.

Wait before migrating a critical workload if you need a mature compatibility matrix, predictable cold-start behavior, or a feature your stack depends on but TorchTPU does not yet expose. Stage the change with one model, one TPU slice, one fixed evaluation set, and separate measurements for installation time, first compilation, steady-state throughput, and output correctness.

TorchTPU FAQ

Is TorchTPU the same as PyTorch/XLA?

No. TorchTPU is a newer native PyTorch backend designed around eager execution and torch.compile; PyTorch/XLA is the older TPU integration path. Existing XLA code may need a separate migration assessment.

Can I install TorchTPU on a laptop without a TPU?

You may install packages, but you cannot validate TPU execution without compatible Google TPU infrastructure. The official repository’s practical setup assumes a TPU VM.

Why did my first TorchTPU step run slowly?

The first execution can compile a graph and populate caches. Measure later iterations separately and keep shapes stable while diagnosing recompilation.

Does changing cuda to tpu guarantee the same performance?

No. The device replacement is a correctness port; profile shapes, sharding, operator coverage, and warm-up before comparing performance.

Is TorchTPU ready for production?

Google describes production inference and training use, but public ecosystem components and features have different maturity levels. Test the exact model, serving engine, and distributed pattern before committing.

Install the pinned wheel on a TPU VM, run the 2×2 verification, and compare one warm workload with your current GPU path. This separates code portability, compiler warm-up, operator coverage, and hardware economics.