Read time: ~12 minutes.

Key facts:

• TorchTPU is a new open-source, native PyTorch backend for Google Cloud TPUs, built by Google with the PyTorch team at Meta. It adds a tpu device to stock PyTorch, so porting starts with .to("cuda") → .to("tpu") and backend="nccl" → backend="tpu_dist" (PyTorch blog, TorchTPU docs).
• The current public release is torch_tpu 0.2.0 on PyPI, with Linux x86_64 wheels for Python 3.11 to 3.14 (PyPI, GitHub README).
• Supported hardware: TPU v5e, v5p, v6e (Trillium) and v7x (Ironwood). TPU v8 support is on the roadmap (FAQ).
• Google’s FAQ states that PyTorch/XLA is deprecated and replaced by TorchTPU. If you have torch_xla code, this is the migration target.
• Google says it validated 4,917 Hugging Face models: 98% (4,821) ran without crashing, and 97% of architecture representatives (638 of 657) matched GPU numerics (PyTorch blog).

What we checked: we read the PyTorch launch post, the TorchTPU README and its docs (VM quickstart, intro, FAQ, feature set, API reference, distributed support, dynamic shapes, numerical precision), the Google Developers post on the new vLLM and SGLang engines, the vLLM-TorchTPU README, and Google Cloud’s TPU VM pages. We confirmed the 0.2.0 wheels on PyPI. We do not have a Cloud TPU, so we did not run these commands ourselves; every command, flag and environment variable below comes from those sources.

If you have a PyTorch codebase and avoided TPUs because PyTorch/XLA meant lazy tensors, marked compile regions and long compile stalls, TorchTPU is the release that changes the trade-off. Google built it as a regular PyTorch backend: eager mode works, torch.compile goes through Dynamo, and collectives plug into torch.distributed. This guide takes you from an empty Google Cloud project to a working TPU tensor, then a compiled model, a multi-chip job and an OpenAI-compatible vLLM server, with the common slowdowns and errors along the way.

If you are on AWS instead, our TorchNeuron guide for Trainium covers the same kind of native PyTorch backend on Amazon’s chips. If you just need a hosted model and not your own accelerator, the Gemini 3.6 Flash API guide is the cheaper starting point. More builder tutorials are on the 7minai homepage.

---

1. What TorchTPU changes (and what it doesn’t)

Before TorchTPU, PyTorch reached TPUs through PyTorch/XLA. The PyTorch blog lists its problems plainly: it “required intrusive source-code modifications to explicitly define compilable regions, suffered from severe compile-time latency, lacked support for PyTorch’s native imperative (‘eager’) execution mode, and lagged in supporting emerging distributed computing primitives” (PyTorch blog).

TorchTPU takes a different route:

	PyTorch/XLA (old path)	TorchTPU (new path)
Device string	xla via torch_xla APIs	tpu, registered automatically on import torch
Eager mode	Lazy tensors	Real eager, plus a fused eager mode
Compile	Lazy-tensor tracing	torch.compile via Dynamo, backend="tpu"
Distributed	XLA-specific APIs	torch.distributed with backend="tpu_dist"; DDP, FSDP2, SimpleFSDP, DTensor
JAX at runtime	Depends on JAX runtime components	No JAX runtime dependency
Status	Deprecated per Google’s FAQ	Active, open source (Apache 2.0)

What it does not do is make GPU-tuned code run at peak TPU speed by itself. Google says so directly: changing .to("cuda") to .to("tpu") should make a model run, but “one cannot achieve peak TPU performance with that simple change, since a model’s shapes, sharding and ops may be skewed toward a GPU architecture.” Plan for a working port first and tuning second.

It is also not a JAX replacement. The FAQ positions it for teams with existing PyTorch code that want GPU/TPU portability or use PyTorch-native frameworks such as TorchTitan, vLLM, SGLang and Hugging Face (FAQ).

---

2. Two ways to try it

Free first look: Colab. The TorchTPU intro and VM quickstart docs each ship as notebooks with an “Open in Colab” badge. The intro notebook says you can run its example on Colab “by enabling the TPU runtime.” Use this to confirm the API feels right before you pay for a VM.

Real work: a Cloud TPU VM. The TorchTPU docs call a TPU VM “the simplest way to use TorchTPU” and suggest a v6e instance as the easiest option (VM quickstart, README). For the newest hardware (Ironwood, TPU7x), the docs point you to GKE instead. Google Cloud’s own docs note that the Cloud TPU API used by gcloud compute tpus tpu-vm is not supported on TPU7x and later.

The rest of this guide uses a single-host v6e VM.

---

3. Create a v6e TPU VM

You need a Google Cloud project with the Cloud TPU API enabled and TPU quota. Google’s PyTorch quickstart lists the IAM roles needed to create and SSH into a TPU: TPU Admin (roles/tpu.admin), Service Account User (roles/iam.serviceAccountUser) and Compute Viewer (roles/compute.viewer) (Google Cloud quickstart).

Set variables. For v6e, Google’s software-version table lists v2-alpha-tpuv6e as the runtime for PyTorch and JAX (TPU software versions). Pick a zone that offers v6e from Google’s TPU regions and zones page.

export PROJECT_ID=your-project-id
export TPU_NAME=torchtpu-test
export ZONE=your-v6e-zone
export ACCELERATOR_TYPE=v6e-1        # or v6e-8 for a full 8-chip host
export RUNTIME_VERSION=v2-alpha-tpuv6e

Which size? Google’s v6e page describes v6e-1 as a single-chip VM “primarily intended for testing” and v6e-8 as a full-host VM with all 8 chips on one VM, optimized for single-workload serving. Each v6e chip has 32 GB of HBM and 918 bf16 TFLOPs peak (TPU v6e).

Create the VM and connect:

gcloud compute tpus tpu-vm create $TPU_NAME \
     --project=$PROJECT_ID \
     --zone=$ZONE \
     --accelerator-type=$ACCELERATOR_TYPE \
     --version=$RUNTIME_VERSION

gcloud compute tpus tpu-vm ssh $TPU_NAME \
    --project=$PROJECT_ID \
    --zone=$ZONE

If SSH fails, the VM may not have an external IP address. Google documents a separate path for VMs without a public IP (Manage TPU resources).

Cost control: a running TPU VM keeps incurring charges, so delete it when you finish:

gcloud compute tpus tpu-vm delete $TPU_NAME --zone=$ZONE --quiet

---

4. Install TorchTPU

The TorchTPU VM quickstart starts with a Python 3.12 virtual environment:

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

Then install the pinned public release. The README’s command adds PyTorch’s CPU wheel index:

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

The --extra-index-url matters. The README warns that installing without it “will install the CUDA version of torch”, which you don’t want on a TPU VM. If you use uv, prepend uv to the pip command.

Python support follows PyTorch’s release matrix; the README currently says Python 3.11+, and PyPI has 0.2.0 wheels for CPython 3.11, 3.12, 3.13, 3.14 and free-threaded 3.14t, all manylinux_2_31_x86_64.

You do not import torch_tpu yourself. According to the API reference, “when torch is imported with TorchTPU installed, the backend automatically registers the "tpu" device and the torch.tpu module” through PyTorch’s autoload mechanism.

---

5. Verify the TPU and run your first tensor

This is the intro notebook’s quickstart:

import torch

device = torch.device("tpu")
x = torch.randn((4, 4), device=device)
y = torch.ones((4, 4), device=device)

result = torch.matmul(x, y).cpu()
print(f"Executed on device: {x.device}")
print(f"Result shape on CPU: {result.shape}, sum: {result.sum():.4f}")

To check what hardware you have, use the device-agnostic torch.accelerator API plus torch.tpu.get_device_properties() for TPU-specific fields (API reference):

import torch

print(torch.accelerator.is_available())
print(torch.accelerator.device_count())
print(torch.accelerator.current_device_index())
print(torch.tpu.get_device_properties().name)   # e.g. "TPU v6e"

get_device_properties() also exposes devices_per_chip, hbm_bytes_per_chip, hbm_bytes_per_device, bf16_flops_per_chip, sparse_cores_per_device and more.

Use torch.accelerator, not torch.tpu.*, for generic calls. On PyTorch 2.11 and later, the redundant helpers such as torch.tpu.device_count(), torch.tpu.is_available() and torch.tpu.synchronize() are internal, and calling them from user code raises an error. The replacements are torch.accelerator.device_count(), torch.accelerator.is_available() and torch.accelerator.synchronize().

Moving data works like CUDA:

cpu_tensor = torch.ones((5, 5))
tpu_tensor = cpu_tensor.to("tpu")                 # move
direct = torch.randn((5, 5), device="tpu")        # create on device
print((tpu_tensor + direct).cpu().sum())          # .cpu() brings results back

---

6. Port an existing training script

The PyTorch blog shows the whole diff for a typical distributed script:

Before (GPU)	After (TPU)
device = torch.device("cuda")	device = torch.device("tpu")
dist.init_process_group(backend="nccl")	dist.init_process_group(backend="tpu_dist")

If you want one script for both, the blog’s third column uses the device-agnostic form:

import torch
import torch.distributed as dist

device = torch.accelerator.current_accelerator()
dist.init_process_group()
model = model.to(device)

Three habits from GPU code need checking:

1. CUDA-only calls. Replace torch.cuda.* with the torch.accelerator equivalents from the API reference (for example, torch.cuda.set_device(rank) becomes torch.accelerator.set_device_index(rank)).
2. FlexAttention. TorchTPU “does not currently have native FlexAttention support.” In Google’s own TorchTitan port, the team swapped FlexAttention for F.scaled_dot_product_attention.
3. CNN memory layout. The docs warn that PyTorch’s default channel-first layout “is very inefficient on TPU” and recommend PyTorch’s channels-last memory format.

For Hugging Face models, Google says its CI runs about 5,000 popular checkpoints (Qwen, Gemma, Llama, Mistral, DeepSeek, Whisper, ViT, FLUX and others) in eager and compiled modes on TPU v5e and newer, so most standard Transformers models should load and run after the device swap. The exception to watch is generation, covered in section 8.

---

7. Choose an eager mode

TorchTPU has three eager modes, set by environment variable or by context manager (intro, PyTorch blog):

Mode	How to enable	Use it for
Strict (default)	nothing	Normal eager runs; ops dispatch one at a time, asynchronously, like CUDA
Debug	export TORCH_TPU_LAUNCH_BLOCKING=1	Finding which op fails; synchronous and slow, like CUDA_LAUNCH_BLOCKING, with bounds checks
DeferAndFuse	export TORCH_TPU_DEFER_AND_FUSE=1	Faster eager; fuses runs of ops until a materialization point such as .cpu() or .item()

Google reports DeferAndFuse “can be 50% or more faster than Strict Eager while reducing peak HBM memory.” In its FLUX TorchTitan case study, switching to DeferAndFuse cut step time from 11.49 s to 9.95 s.

To scope a mode to one block instead of the whole process:

with torch.tpu.set_eager_mode(torch.tpu.EagerMode.DEFER_AND_FUSE):
    x = torch.ones((4, 4), device="tpu") + 1
    print(torch.tpu.eager_mode, x.sum().item())

torch.tpu.EagerMode and set_eager_mode are listed as experimental APIs. They may change without a deprecation cycle and emit a one-time warning when used (public APIs).

---

8. Compile with torch.compile

For production speed, compile. TorchTPU registers a Dynamo backend:

compiled_model = torch.compile(model, backend="tpu")

Dynamo hands TorchTPU the FX graph, and TorchTPU lowers the whole graph to one StableHLO program for the XLA compiler. In the FLUX case study, full torch.compile(backend="tpu") took step time from 5.55 s to 3.02 s, the largest single gain in Google’s optimization sequence.

Dynamic shapes are the main trap. The PyTorch blog says that, as a temporary limitation, TorchTPU automatically sets dynamic=False when you call torch.compile on a TPU machine. The intro doc adds that dynamic=None and dynamic=True “are not yet supported for production.” In practice:

• Every new input shape triggers an XLA recompile.
• The docs give a concrete example: “running decode/generate step on Hugging Face LLMs will involve recompilations at each step unless bounded dynamism is enabled.”

The fix in progress is bounded dynamism. You declare a range for a dimension with torch._dynamo.mark_dynamic(tensor, dim, min=..., max=...), and TorchTPU compiles a small number of padded buckets (bucket size defaults to 128, up to 8 buckets) instead of one graph per exact shape (dynamic shapes doc). The feature table marks SymInts as experimental. Until it is production-ready, the reliable approach is the classic TPU one: pad batches and sequence lengths to a few fixed sizes. For LLM text generation, use vLLM-TorchTPU (section 11), which already handles this.

---

9. Turn on the persistent compilation cache

Both eager warmup and torch.compile spend time compiling XLA programs. The docs’ recommendation for both problems is the same: “Enable persistent compilation cache.”

TorchTPU has three cache tiers (API reference):

• Tier 1: in-process memory, always on.
• Tier 2: host-local shared memory, so several processes on one VM reuse compiles. Set TORCH_TPU_TIER2_COMPILATION_CACHE to a directory name (not a path); it is placed under /dev/shm/torch_tpu_cache/<name>.
• Tier 3: persistent across runs and hosts. Set TORCH_TPU_TIER3_COMPILATION_CACHE_ROOT to a gs:// bucket URI or a shared filesystem path.

The API reference sets them before importing torch:

import os

os.environ["TORCH_TPU_TIER2_COMPILATION_CACHE"] = "tpu_tier2_cache"
os.environ["TORCH_TPU_TIER3_COMPILATION_CACHE_ROOT"] = "/path/to/shared/cache"  # or gs://your-bucket/cache

import torch
a = torch.randn(1024, 1024, device="tpu")
_ = (a @ a).to("cpu")

Compiled executables are written asynchronously, so the second run of the same script skips most compilation.

---

10. Run on all chips of a host with torchrun

TorchTPU uses PyTorch’s one process per device model. One difference from GPUs: a TPU process sees only its own device, so the runtime needs the slice topology before init_process_group. Two environment variables carry it, TORCH_TPU_TOPOLOGY (for example "2,2,1") and TORCH_TPU_SLICEBUILDER_ADDRESSES (one host:port per device). On a single host or a GKE TPU slice, TorchTPU can fill them in for you (distributed support).

Worker script, train.py, from the docs:

import torch
import torch.distributed as dist

def worker_main():
    dist.init_process_group(backend="tpu_dist")
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    device = torch.device("tpu")
    print(f"[Rank {rank}/{world_size}] Initialized on {device}")

    x = torch.ones(4, 4, device=device) * (rank + 1)
    dist.all_reduce(x, op=dist.ReduceOp.SUM)
    print(f"[Rank {rank}] Reduced sum: {x[0, 0].item()}")

    dist.destroy_process_group()

if __name__ == "__main__":
    worker_main()

Launch it:

# Detects and exports TORCH_TPU_TOPOLOGY, TORCH_TPU_SLICEBUILDER_ADDRESSES,
# WORLD_SIZE, NODE_RANK and multi-host coordination variables
eval $(python3 -m torch.tpu.distributed | sed 's/^/export /')
torchrun --nproc_per_node=$WORLD_SIZE train.py

If you launch workers from Python instead of torchrun, call torch.tpu.distributed.set_tpu_launch_env() in the parent process before spawning. Both methods leave variables you already set untouched, so you can override individual values.

Two scaling notes from Google:

• Device count is not always chip count. In the blog’s Muse Glimmer 30B example, a single TPUv7 chip “exposes two devices to PyTorch (corresponding to the two included TensorCores),” so one chip ran with a world size of 2. Use WORLD_SIZE from the helper rather than hard-coding chips.
• FSDP choice matters. The docs recommend SimpleFSDP with torch.compile over FSDP2, which the feature table lists as experimental and “low perf.” For mesh layout, torch.tpu.distributed.topology_aware_mesh() maps your most communication-heavy parallel dimension (usually tensor parallel) onto the fastest physical links.

By default, TorchTPU compiles each collective in its own small module so ranks stay consistent. That is safe but, per the blog, can cost “30-40% end-to-end” because compute and communication can’t overlap. If a region is identical on every rank, the experimental @torch.tpu.distributed.spmd_safe decorator removes those graph breaks. Only use it when no code inside branches on get_rank().

---

11. Serve an LLM with vLLM-TorchTPU

For LLM inference you don’t have to write the decode loop yourself. Google and Inferact (founded by vLLM’s creators) released vLLM-TorchTPU, an out-of-tree vLLM platform plugin that runs vLLM’s standard PyTorch model definitions on TPU with Pallas kernels. An SGLang equivalent, SGLang-TorchTPU, was built with RadixArk (Google Developers Blog).

Install it on a TPU VM with Python 3.12 (vLLM-TorchTPU README):

pip install \
  --extra-index-url https://download.pytorch.org/whl/cpu \
  vllm-torchtpu

Or use the container: docker pull vllm/vllm-torchtpu:latest.

Serve and query it exactly like any vLLM server:

vllm serve meta-llama/Llama-3.2-1B-Instruct \
  --tensor-parallel-size 1 \
  --max-model-len 4096

curl http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "meta-llama/Llama-3.2-1B-Instruct",
    "messages": [{"role": "user", "content": "Explain how TPUs accelerate matrix multiplication."}],
    "max_tokens": 128
  }'

Things to know before you rely on it:

• Hardware status. The README marks TPU v7x (Ironwood) as recommended and CI-tested. v6e, v5e and v5p are experimental and not covered by CI. On the v6e VM from this guide, expect rough edges.
• First start is slow. The server compiles TPU graphs for the configured shapes, which “can take several minutes.” Set VLLM_XLA_CACHE_PATH to keep compiled artifacts across restarts.
• Model coverage. The README lists first-class support for Qwen3/Qwen3.5, DeepSeek-V3/R1/V4, Kimi-K2.5/K3, GLM-5, Llama 3/4, GPT-OSS, Gemma 3/4, Mixtral and vision-language families.
• The JAX engine is not going away. The old tpu-inference backend is now called vLLM JAX and stays maintained. vLLM-TorchTPU is the recommended path for PyTorch model definitions.

For general vLLM deployment choices (tensor parallelism, context length, FP8), our Qwen3.8-Flash-Next vLLM guide applies on the serving side, even though it targets GPUs.

---

12. Numerical precision: why your float32 results differ slightly

On CUDA, PyTorch’s default float32 matmul precision is highest. On TPU, TorchTPU defaults to medium: matmul inputs are converted to bfloat16 for the MXU, with float32 accumulation (PyTorch blog). Google says this is accurate enough for “nearly all deep learning workloads,” but small differences from GPU outputs are expected.

If a specific layer needs more precision, scope it instead of changing the whole run:

a = torch.randn(64, 64, device="tpu", dtype=torch.float32)
b = torch.randn(64, 64, device="tpu", dtype=torch.float32)
with torch.tpu.precision(torch.tpu.Precision.HIGH):
    c = torch.matmul(a, b)

The costs are real: high uses a 3-pass algorithm at roughly 3× MXU compute, and highest uses 6 passes at roughly 6×. TorchTPU emits a one-time performance warning when you select either. The global setting is still torch.set_float32_matmul_precision.

One related flag: torch.tpu.allow_excess_precision (default False) lets XLA keep fused intermediates in float32. Set it at the very start of the script. The docs say the runtime locks it at the first compilation, and changing it afterward raises a RuntimeError.

---

13. Troubleshooting

pip downloads a huge CUDA build of torch. You installed without --extra-index-url https://download.pytorch.org/whl/cpu. Reinstall in a clean venv using the command in section 4.

Calling torch.tpu.device_count() (or is_available, synchronize) raises an error. On PyTorch 2.11+, these are internal. Use torch.accelerator.device_count(), torch.accelerator.is_available() and torch.accelerator.synchronize().

A cryptic RuntimeError from deep in the backend. Set TORCH_SHOW_CPP_STACKTRACES=1 before importing torch to get the C++ traceback. To find which op fails, set TORCH_TPU_LAUNCH_BLOCKING=1 so ops run synchronously.

Silent wrong results from indexing ops. TPU skips bounds checks by default for speed. TORCH_TPU_ENABLE_DEBUG_CHECKS=1 turns them on, so an out-of-range torch.take raises take(): expected indices to be in range instead of corrupting data. It slows execution, so use it only while debugging.

Every step is slow, especially in a generate loop. Changing shapes trigger recompiles. Pad to fixed sizes, enable the tier-2/3 compilation cache (section 9), and for LLM serving use vLLM-TorchTPU.

A second process can’t open the TPU. Only one OS process can access a given TPU device. That is why TorchTPU uses one process per device, and why co-located RL setups (trainer and sampler sharing a device) aren’t supported out of the box.

Multi-chip training is much slower than expected. Check that you use SimpleFSDP with torch.compile rather than FSDP2, that your mesh comes from topology_aware_mesh(), and whether a rank-identical region can use @spmd_safe.

Changing allow_excess_precision raises RuntimeError. Set it before any TPU op runs.

---

FAQ

Is TorchTPU free? The software is open source under Apache 2.0. You pay for the Cloud TPU VM or GKE nodes it runs on. Colab’s TPU runtime is the no-setup way to try the notebooks.

Should I migrate from PyTorch/XLA now? Google’s FAQ says PyTorch/XLA “is deprecated and replaced by TorchTPU.” New TPU projects in PyTorch should start on TorchTPU. Existing torch_xla code needs a port: device strings, the process-group backend and any xm.* calls.

Which TPU should I use? For a first port, Google’s v6e page describes the single-chip v6e-1 VM as intended for testing. For LLM serving with vLLM-TorchTPU, v7x (Ironwood) is the only CI-tested generation. Ironwood isn’t available through the Cloud TPU API, and the TorchTPU docs use GKE for recent hardware.

Can I write custom TPU kernels from PyTorch? Yes. Pallas kernels wrapped with torch.tpu.pallas.jax_op and the Tokamax kernel library are listed as production-ready; Helion is experimental. Note that torch.tpu.pallas.jax_op itself is in the experimental API list.

Does it support quantization? The feature table lists full-tensor FP8/FP4 and per-channel FP8 as validated. Block-wise FP8/FP4 with torchao is on the near-term roadmap.

Is there tooling to help with onboarding? The FAQ points to the experimental TPU agent suite, an open-source plugin with sub-agents for onboarding, diagnostics and inference tuning on v6e and v7x.

---

Sources

• Introducing TorchTPU (PyTorch blog)
• google-pytorch/torch_tpu on GitHub and its docs: VM quickstart, intro, FAQ, feature set, API reference, public APIs, distributed support, dynamic shapes
• torch-tpu 0.2.0 on PyPI
• Announcing new vLLM and SGLang inference engines on TPU (Google Developers Blog)
• vllm-project/vllm-torchtpu README
• Google Cloud docs: PyTorch quickstart, TPU software versions, TPU v6e, Manage TPU resources