TorchTPU Documentation: Native PyTorch Backend for Google TPUs
by Google and Meta (PyTorch)
Official docs for running existing PyTorch training and inference code on Google Cloud TPUs by changing the device to "tpu".
Overview
TorchTPU is a native TPU backend for PyTorch that Google started in mid 2025 with the PyTorch team at Meta and announced on the PyTorch blog on October 20, 2026. The documentation (version 0.2.0, matching the latest PyPI release) is organized into Get started (introduction, getting started on a TPU VM or on GKE, the supported feature set and roadmap, an API reference and an FAQ), In-depth architecture and core concepts (how TorchTPU works, the execution model, the memory model, distributed support, custom kernel support and dynamic shapes), Case studies (FLUX optimization, Muse Glimmer 30B inference and training, and using FSDP) and Analyze your workload (numerical precision and parity, and profiling with XProf). The pages are built from ready-to-run notebooks. Unlike PyTorch/XLA, TorchTPU has no JAX runtime dependency and treats eager execution as the default: the backend registers as a PrivateUse1 device, lowers to StableHLO, adds a DeferAndFuse eager mode, plugs into torch.compile through a Dynamo backend, and supports DDP, FSDP and DTensor through TPU collectives. Custom kernels can be written in JAX/Pallas or in PyTorch's Helion DSL. According to the launch post, about 450 ATen ops are covered and 4,821 of 4,917 popular Hugging Face models ran without crashing, with numerical parity against GPU on 638 of 657 architecture representatives.
At a Glance
- Topic
- Frameworks
- Level
- Intermediate
- Format
- Documentation
- Cost
- Free
- Duration
- Self-paced reference with runnable Marimo notebooks; running them needs a Cloud TPU VM (v6e is the easiest per the README)
- Provider
- Google and Meta (PyTorch)
- Hands-on
- Yes — code/exercises
- Certificate
- None
What You’ll Learn
- ✓Install the torch_tpu wheel from PyPI against the CPU PyTorch index on a TPU VM
- ✓Port an existing PyTorch model to TPU by moving modules and tensors to the tpu device
- ✓Choose between eager execution, DeferAndFuse and torch.compile for a given TPU workload
- ✓Shard training across TPU chips with FSDP, DDP and DTensor device meshes mapped to ICI topology
- ✓Write custom TPU kernels in JAX/Pallas or Helion and call them like PyTorch functions
- ✓Handle dynamic shapes with bounded dynamism to avoid repeated XLA recompilation
- ✓Check numerical parity against GPU baselines and profile TPU performance with XProf
- ✓Deploy TorchTPU workloads on Google Kubernetes Engine as well as on single TPU VMs
Highlights
- •Built by Google with Meta's PyTorch team and integrated into Hugging Face Transformers v5, Diffusers and Kernels, plus vLLM, SGLang, TorchTitan, Ray and TorchRec
- •The launch post reports a pure-PyTorch Muse Glimmer 30B reaching about 25.7k tokens/s prefill at 59% MFU on a single TPUv7 chip with no model-specific kernels
- •Open source under Apache 2.0 at google-pytorch/torch_tpu, which had 804 stars and 77 forks when checked on October 31, 2026
- •Covers the parts PyTorch/XLA users struggled with: eager mode, compile latency via a compilation cache, and native distributed idioms
- •Still early (0.2.0): the docs include a supported feature set and roadmap page, so check it before committing a production workload
Who It’s For
Best For
- ✓PyTorch engineers who have Cloud TPU capacity and do not want to rewrite models in JAX
- ✓Teams moving training or fine-tuning jobs from PyTorch/XLA to a native TPU backend
- ✓Inference engineers serving open-weight models on TPUs with vLLM or SGLang
- ✓ML platform teams comparing TPU and GPU cost for large PyTorch workloads
Prerequisites
- •Solid PyTorch experience, including torch.compile and torch.distributed basics
- •A Google Cloud account with TPU quota (TPU VMs or GKE) to run the notebooks
- •Python 3.11 or newer; Bazel (via Bazelisk) only if you build the wheel from source
- •Helpful but optional: familiarity with XLA, StableHLO or Pallas for kernel and performance work
FAQ
What is TorchTPU Documentation: Native PyTorch Backend for Google TPUs?
The official documentation for TorchTPU, the open-source PyTorch backend for Google TPUs that Google and Meta announced in October 2026. It is for PyTorch engineers who want to train, fine-tune or serve models on Cloud TPUs with eager mode, torch.compile and native distributed APIs instead of rewriting code for JAX or PyTorch/XLA.
Is TorchTPU Documentation: Native PyTorch Backend for Google TPUs free?
TorchTPU Documentation: Native PyTorch Backend for Google TPUs is free to access.
What level is TorchTPU Documentation: Native PyTorch Backend for Google TPUs for?
TorchTPU Documentation: Native PyTorch Backend for Google TPUs is aimed at a intermediate audience. Recommended background: Solid PyTorch experience, including torch.compile and torch.distributed basics, A Google Cloud account with TPU quota (TPU VMs or GKE) to run the notebooks, Python 3.11 or newer; Bazel (via Bazelisk) only if you build the wheel from source, Helpful but optional: familiarity with XLA, StableHLO or Pallas for kernel and performance work.
How long does TorchTPU Documentation: Native PyTorch Backend for Google TPUs take?
Expect roughly Self-paced reference with runnable Marimo notebooks; running them needs a Cloud TPU VM (v6e is the easiest per the README). Most learners work through it at their own pace.
What will I learn from TorchTPU Documentation: Native PyTorch Backend for Google TPUs?
You'll learn: Install the torch_tpu wheel from PyPI against the CPU PyTorch index on a TPU VM; Port an existing PyTorch model to TPU by moving modules and tensors to the tpu device; Choose between eager execution, DeferAndFuse and torch.compile for a given TPU workload; Shard training across TPU chips with FSDP, DDP and DTensor device meshes mapped to ICI topology; Write custom TPU kernels in JAX/Pallas or Helion and call them like PyTorch functions; Handle dynamic shapes with bounded dynamism to avoid repeated XLA recompilation; Check numerical parity against GPU baselines and profile TPU performance with XProf; Deploy TorchTPU workloads on Google Kubernetes Engine as well as on single TPU VMs.
Topics
Sources
This page was written from 3 sources, 2 on domains other than torchtpu.readthedocs.io.