Skip to content
AI Coding

How to Use TorchTPU: Native PyTorch on TPUs Guide

Google open-sourced TorchTPU, a native PyTorch backend for TPUs. Hands-on TPU VM setup, Strict vs DeferAndFuse vs compile, and the gotchas that break the one-line swap story.

6 min readBeginner

Everyone’s celebrating “just change .to(‘cuda’) to .to(‘tpu’)” like the hardware wars are over. They’re not. TorchTPU makes TPUs feel like real PyTorch devices for the first time, but the silicon still has opinions about shapes, processes, and precision that GPU code never had to care about. Treat the one-line swap as magic and you get correct-but-slow runs – then wonder why launch-blog MFU numbers never show up on your job.

You’re a PyTorch person tired of CUDA version soup or GPU quota queues. On October 20, 2026, at the PyTorch Conference, Google open-sourced TorchTPU with Meta’s PyTorch team. This is the hands-on path: stand it up on a TPU VM, pick a mode that matches your stage, and read the sharp edges before you migrate a real training or serving job.

What TorchTPU Actually Changes

PrivateUse1 registration. That’s the whole trick. TorchTPU plugs in as a native device backend – no JAX runtime under the hood – so tensors stay plain torch.Tensor objects. You point at torch.device("tpu") and the dispatcher sends ATen ops to TPU kernels. Coverage sits around 450 ops as of the v0.2.0 launch docs, enough for mainstream LLM and vision stacks, not every custom op under the sun.

Old PyTorch/XLA lazy tensors are the thing this replaces. Eager is default, so breakpoints and print behave closer to CUDA. Production path is torch.compile(model, backend="tpu"): TorchDynamo → StableHLO → XLA. Multi-chip uses the normal process-group API with backend="tpu_dist".

Hooks landed with the launch posts – not a roadmap slide. Hugging Face Transformers v5 takes device_map="tpu". There are TorchTPU builds/plugins for vLLM and SGLang, plus TorchTitan and Ray paths. Full write-up lives on the Google Cloud announcement; the architecture notes are on the PyTorch blog.

Practical Setup on a TPU VM

Laptop won’t cut it. You need a Google Cloud TPU VM. README energy still points at v6e as the least painful first box (as of the Oct 2026 docs – confirm the current SKU before you click create). gcloud create, SSH, then:

python -m venv .venv
source .venv/bin/activate
pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpu

Skip --extra-index-url and pip silently grabs the CUDA torch wheel. On a TPU VM that wheel is huge, wrong, and a classic source of import/dispatcher ghosts. Python 3.11+ only. Pin notes and troubleshooting sit in the google-pytorch/torch_tpu repo and torchtpu.readthedocs.io.

Smoke test:

import torch
device = torch.device("tpu")
x = torch.randn(4, 4, device=device)
y = torch.ones(4, 4, device=device)
print((x @ y).cpu().sum())

Number prints, device string says tpu – you’re live. Multi-chip: set the launch env (or torch.tpu.distributed.set_tpu_launch_env()), then torchrun --nproc_per_node=N your_script.py with dist.init_process_group(backend="tpu_dist").

Pro tip: fresh venv. CUDA torch leftovers mixed with torch_tpu is the fastest route to mysterious failures.

Modes That Actually Matter

Strict eager is the default: one op at a time, async, CUDA-shaped mental model. Correct. Debuggable. Also a tax – TPUs hate a storm of tiny kernels.

  • Debug – export TORCH_TPU_LAUNCH_BLOCKING=1. Sync launch plus bounds checks. Flip this when results look “fine” but wrong.
  • DeferAndFuse – export TORCH_TPU_DEFER_AND_FUSE=1. Batches work until a materialization point (.cpu(), .item(), etc.). Launch docs put it at 50%+ faster than Strict with lower peak HBM.
  • Compile – model = torch.compile(model, backend="tpu"). This is where big step-time drops usually come from. Cache helps across restarts. Today compile still forces dynamic=False; change batch/seq freely and you recompile unless you pad or wait on bounded dynamism (still maturing per the same launch notes).

Precision default is “medium”: BF16 matmuls on the MXU, FP32 accumulate. That matches how the chip wants to run. high / highest exist for stubborn convergence bugs – and they cost roughly 3× / 6× on matmul compute. Use the context manager on one sensitive layer instead of flipping the whole model.

Serving drop-in: vLLM-TorchTPU or SGLang-TorchTPU, keep OpenAI-compatible front ends. Training recipes in TorchTitan reported up to 52% MFU on Qwen 3 8B with near-linear scale out to 8,192 Ironwood (TPU v7x) chips as of the Oct 20, 2026 posts – treat that as a tuned ceiling, not your first Strict run.

Is the “native” story done? For stock LLM/vision graphs, ops + compile are usable. Custom attention twists and heavy dynamic control flow still ask for extra kernel or shape work. Hype slides skip that sentence.

Honest Limitations (Read Before You Migrate)

Software: Apache 2.0. Silicon: still Google Cloud. Quotas and price still gate long runs – check both before a multi-day job.

The catch is process topology. One user-space process owns one TPU device. Trainer+sampler colocation that some GPU RL loops assume fails out of the box. Disaggregate with Ray or a VeRL-style TPU backend.

Shapes are not a GPU afterthought. Wide dimensions that feed the MXU well win; awkward aspect ratios waste systolic cycles. Pair that with the compile dynamic=False constraint above and you either pad or eat recompiles.

Out-of-box Strict can look awful on real fine-tunes. One published ladder (MarkTechPost summarizing the official case) went from 11.49 s/step down to 2.78 s after DeferAndFuse + compile + kernel/shape/sharding work. The device swap gets you correctness. The ladder gets you TCO.

Bounds checks are off by default for speed. Out-of-range indexing can corrupt quietly instead of throwing. When you smell it, turn on TORCH_TPU_LAUNCH_BLOCKING=1 / debug checks before you blame the optimizer.

FAQ

Do I still need PyTorch/XLA?

No. New work should start on TorchTPU. XLA is the predecessor path.

Can I keep the same training script and just flip the device string?

For most pure PyTorch scripts: change device + tpu_dist, run single-chip, check numerics. Peak speed still wants DeferAndFuse or compile, shape hygiene, and sometimes a hot-spot kernel. Then scale.

Should I move off GPUs right now?

If you already live in Hugging Face + vLLM/SGLang and GPU cost or quota is the pain, spin a small TPU VM this week and time your model under Strict, then DeferAndFuse. That one comparison beats another launch thread. If the codebase is CUDA graphs, custom Triton, and multi-stream in-place tricks, wait while more kernels and agent tooling land – the “always GPU” default is no longer automatic, but it is not obsolete overnight.

Next action: provision a v6e TPU VM, run the exact pip line above, time one existing model Strict vs DeferAndFuse. That single A/B tells you more than the keynote.

Related Tutorials