Until October 20, 2026, getting a PyTorch model onto a TPU meant going through PyTorch/XLA or rewriting it in JAX. That day, at PyTorch Conference 2026, Google open-sourced TorchTPU, a native PyTorch backend built together with Meta's PyTorch team. The TPU becomes just another device, like cuda or cpu, and Hugging Face already supports it in Transformers v5 and Diffusers.
If you train or serve models on Google Cloud, this matters more than the announcement suggests: the question "is it worth porting to TPU?" no longer starts with "how much code do we have to touch?".
What changes compared with PyTorch/XLA
PyTorch/XLA (torch_xla) worked, but it came with its own mental model. Hugging Face sums it up well in its announcement: tensors were lazy, so code didn't always run when you expected, small changes triggered graph recompilations, and debugging was hard to reason about. To bridge the gap, Hugging Face maintained Optimum TPU, which TorchTPU now makes unnecessary for running Transformers and Diffusers.
TorchTPU flips the approach:
| PyTorch/XLA | TorchTPU | |
|---|---|---|
| Execution model | Lazy tensors and an XLA graph | Eager by default, torch.compile for production |
| Device in code | Its own torch_xla API | "tpu", like any other device |
| Debugging | Recompilations that are hard to follow | Eager mode and standard PyTorch diagnostics |
| Hugging Face | Through Optimum TPU | Native in Transformers v5 and Diffusers |
The idea is to debug in eager mode and switch to torch.compile when you need to squeeze the TPU. The backend ships with Pallas kernels tuned for TPUs and native distributed primitives.
The ecosystem that lands the same day
What makes the launch credible is that it doesn't come alone. Google announced integrations with the engines most teams already run:
- vLLM, with Inferact: a native backend with continuous batching, quantization and multi-chip parallelism, without changing workers, schedulers or API servers.
- SGLang, with RadixArk, for agentic and structured workloads that reuse the KV cache across requests.
- TorchTitan for distributed training. Google reports up to 52% MFU (Model Flops Utilization) on Qwen 3 8B, with near-linear scaling up to 8,192 Ironwood TPUs.
- VeRL (in preview) for reinforcement learning and Ray to orchestrate CPUs, GPUs and TPUs in a single workflow.
- llm-d 1.0 to route inference on Kubernetes across TPUs and GPUs behind a common endpoint on GKE.
Hugging Face adds its own reality check: 96.5% of the Transformers tests that run on TPU pass, and 19 of the 34 issues they have filed on the TorchTPU repo are already fixed. It's an honest number: not everything is green, and they say so.
How to install it and try it
On a TPU VM, installation is a single pip command. The extra index pulls the CPU build of PyTorch; without it, pip downloads the much larger CUDA build, which is useless on a TPU:
pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpuWith Transformers v5, loading a model on the TPU is a matter of passing device_map="tpu":
from transformers import AutoModelForCausalLM, AutoTokenizer
model_id = "Qwen/Qwen3.5-9B"
model = AutoModelForCausalLM.from_pretrained(model_id, device_map="tpu")
tokenizer = AutoTokenizer.from_pretrained(model_id)
inputs = tokenizer("The future of AI is", return_tensors="pt").to("tpu")
out = model.generate(**inputs, max_new_tokens=64)
print(tokenizer.decode(out[0], skip_special_tokens=True))In Diffusers, pipe.to("tpu") is enough for eager mode, and a compile(backend="tpu") call on the transformer gives you compiled mode. In Hugging Face's examples on a TPU v6e-4, a 1024×1024 image with 40 steps drops from 8.0 s in eager mode to 3.5 s compiled. The first call compiles; later calls with the same size reuse the graph.
Checklist before production
If you already run workloads on torch_xla or Optimum TPU, this is what I'd validate before switching:
- Numerical parity. Run the same batch on your current stack and on TorchTPU in eager mode, and compare outputs with a sensible tolerance for
bfloat16. - Recompilations. In compiled mode, pin input shapes (
dynamic=Falsein the Diffusers examples) and count how often it compiles under real traffic; the 8.0 s to 3.5 s gain only exists if the graph is reused. - Serving engine. If you serve with vLLM or SGLang, test their TorchTPU backend with your specific model before assuming it inherits every feature of the GPU version.
- Open issues. Check the TorchTPU issues that touch your architecture: with 96.5% of tests passing, the remaining 3.5% might be exactly your model.
What it means
Until now, choosing TPUs also meant choosing a stack (JAX or XLA). With TorchTPU, the decision goes back to hardware price and availability, which is how it should be. Meta's quote in the announcement points the same way: they don't want to depend on a single architecture. For teams that only run on NVIDIA GPUs, TPUs make the shortlist with no rewrite cost; what remains to be proven is whether TorchTitan's 52% MFU holds up outside Google's own benchmarks.
Sources
- TorchTPU is now open source: Run native PyTorch on TPUs, Google Cloud Blog
- Hugging Face ❤️ TorchTPU: TPUs, Natively, Hugging Face