TorchTPU is Google's open-source PyTorch backend for Cloud TPUs, built with Meta's PyTorch team and open-sourced on 20 October 2026. It registers the TPU as a normal PyTorch device called tpu, so for most code the port starts with device="cuda" becoming device="tpu". It replaces PyTorch/XLA, which Google now calls deprecated.
The library is free under Apache 2.0, and you pay for TPU time: $12.00 per Ironwood chip-hour on demand, $6.00 with flex-start, and $2.70 for Trillium. Google reports 98% of nearly 5,000 Hugging Face models run out of the box. The catch: it's version 0.2.0, torch.compile still forces static shapes, and Google says plainly that "runs" and "runs fast" are two different amounts of work.
If you train or serve your own PyTorch models, this is the first time TPUs are worth a pilot without learning JAX. If you just want AI doing a job, like answering tickets, skip the chip question: an eesel AI teammate handles the model layer, and SamCart runs four of them on Freshdesk and website chat.
What is TorchTPU?
TorchTPU is a native PyTorch backend for Google Cloud TPUs, the custom AI chips Google uses to train and serve Gemini. Google started the project in mid 2025 with the PyTorch team at Meta, described it on the Google Developers Blog in April 2026, and open-sourced it at PyTorch Conference 2026 on 20 October. The code lives at google-pytorch/torch_tpu under Apache 2.0, and the current release is 0.2.0 on PyPI.
I build features and AI agents at eesel, and most of my week is spent on the part of the stack above the chip. So when a launch says "change one line," I go and read how that line actually gets executed. The short version here is good news with an asterisk on it, the line really does get your model running, and the docs are unusually honest about the work which comes after.
The Google Cloud announcement frames the goal simply: TPUs join GPUs and CPUs as a native option for PyTorch. Meta is the headline partner, and its PyTorch engineering director explains why in the launch post:
"We cannot be constrained by a single hardware architecture."
Jana van Greunen, Director of PyTorch Engineering at Meta, on the Google Cloud blog
When I captured the repo on 28 October it had 782 stars, 105 contributors and 4,677 commits, with commits landing hourly. Per the PyTorch blog, it is already "used in production for both inference and training workloads for some of the largest Google Cloud customers."
What's actually new: goodbye PyTorch/XLA
PyTorch has been running on TPUs since 2020 through PyTorch/XLA, so it's a fair question what changed. Google's own answer in the PyTorch blog is blunt: PyTorch/XLA used lazy tensors, needed intrusive code changes to mark compilable regions, had severe compile-time latency and had no real eager mode. The TorchTPU FAQ now says PyTorch/XLA is deprecated and replaced by TorchTPU.
People who lived with the old stack remember it the same way. From the Hacker News thread on the April preview:
"I did trained some research models using the existing PyTorch/XLA on TPUs, and it was a mess of undocumented behavior and bugs (silently hanging after 8 hours of training!)."
TorchTPU plugs into PrivateUse1, PyTorch's official slot for out-of-tree hardware, so a TPU tensor is an ordinary PyTorch tensor. It has no JAX runtime dependency, follows PyTorch's one-process-per-device model, and handles code where rank 0 does a little extra work, which the old SPMD-only stack couldn't.
The other notable change is the ecosystem around it. Google shipped TorchTPU together with backends for vLLM (with Inferact), SGLang (with RadixArk), TorchTitan, Ray (with Anyscale), TorchRec, and native support in Hugging Face Transformers v5 and Diffusers. Hugging Face says its Optimum TPU wrapper is no longer needed.
The timing is worth to notice too. AWS launched TorchNeuron, its native PyTorch backend for Trainium, the same day at the same conference. Both clouds arrived on the same conclusion, PyTorch is where the code lives, so the chip has to look like a PyTorch device.
How TorchTPU works
Under the hood, TorchTPU has two paths, both ending in Google's XLA compiler. The eager path is a C++ backend registered with the PyTorch dispatcher. The compiled path is a TorchDynamo backend that turns a whole FX graph into one StableHLO program, instead of going through TorchInductor like GPUs do.
Eager mode comes in three settings, per the TorchTPU intro docs:
- Strict (default) sends one ATen op at a time to the TPU, asynchronously, like PyTorch on a GPU.
- Debug (
TORCH_TPU_LAUNCH_BLOCKING=1) runs each op synchronously with bounds checks, the TPU version ofCUDA_LAUNCH_BLOCKING. - DeferAndFuse (
TORCH_TPU_DEFER_AND_FUSE=1) holds ops until something like.cpu()or.item()needs a value, then fuses and compiles them together. Google says it can be 50% or more faster than Strict and uses less HBM memory.
All three share a compilation cache with three tiers: per process, shared across processes on one host, and an optional persistent tier on Cloud Storage so a fleet doesn't recompile the same graphs after every restart.
torch.compile(model, backend="tpu") is the production path. The trade is strictness: TorchTPU currently sets dynamic=False for you, because XLA wants static shapes. Google is building "bounded dynamism," where you give a maximum size and the graph compiles once. In early tests it got correct results on 834 of 839 causal LLMs, but the feature table still lists it as experimental.
There's also a precision difference worth knowing before you compare loss curves. TPU matrix units multiply in bfloat16 and accumulate in FP32, so TorchTPU defaults to "medium" precision, while PyTorch on CUDA defaults to full FP32. You can ask for "high" or "highest," but they cost roughly 3x and 6x the matrix compute.
For custom kernels, TorchTPU supports Pallas (Google's kernel language from JAX) and the Tokamax kernel library as production-ready, plus Helion, PyTorch's portable kernel DSL, as experimental. Custom CUDA kernels don't carry over.
To see how much each layer matters, the PyTorch blog walks a FLUX image model through the stack step by step:
To me that chart is the honest version of the "one line" pitch. The one line gets you the 11.49-second step. The 2.78-second step takes five more rounds of work, and torch.compile alone is the biggest single jump.
What the code change looks like
The getting started guide runs on a TPU VM (a v6e is the easiest), with Python 3.11 or newer:
pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpu
Keep the extra index: without it pip pulls the CUDA build of PyTorch, which is much bigger and useless on a TPU. Then the device change:
import torch
device = torch.device("tpu") # was "cuda"
model = MyModel().to(device)
compiled = torch.compile(model, backend="tpu") # optional, for speed
For distributed jobs there's a second edit. dist.init_process_group(backend="nccl") becomes backend="tpu_dist", and TorchTPU needs two TPU-specific settings (the slice topology and the chip addresses) that plain torchrun doesn't know. One command prints them for you:
eval $(python -m torch.tpu.distributed | sed 's/^/export /')
torchrun --nproc_per_node=8 train.py
For Hugging Face models it's shorter still. Per the Hugging Face launch post, AutoModelForCausalLM.from_pretrained(model_id, device_map="tpu") loads weights straight onto the TPU, and a Diffusers pipeline needs only pipe.to("tpu").
That means most open models I've covered, like Qwen, Gemma 4 or DeepSeek V4 Flash, are a device_map away from a TPU.
Google also shipped a TPU agent suite in preview: skills that install into coding agents (Claude Code and Antigravity today, Codex and Cursor "coming soon") to handle onboarding, profiling and inference tuning. It's a sign of where this kind of port is heading. A mechanical swap like this is exactly what Claude Code or OpenAI Codex handle well, as long as you read the compile logs yourself.
How much does TorchTPU cost?
The library is free. TorchTPU is Apache 2.0 and installs from PyPI. What you pay for is TPU time, which Google prices per chip-hour on its Cloud TPU pricing page. Charges run while a TPU node is in a READY state, so an idle node you forgot about still bills.
These are the US rates for the four generations TorchTPU supports, checked on 28 October 2026:
| TPU | Region | On demand | DWS Flex-start | 1-year commit | 3-year commit |
|---|---|---|---|---|---|
| Ironwood (v7) | us-central1 | $12.00 | $6.00 | $8.40 | $5.40 |
| Trillium (v6e) | us-east1 | $2.70 | $1.35 | $1.89 | $1.22 |
| TPU v5p | us-east5 | $4.20 | $2.10 | $2.94 | $1.89 |
| TPU v5e | us-central1 | $1.20 | $0.60 | $0.84 | $0.54 |
Flex-start (part of Google's Dynamic Workload Scheduler) is half the on-demand rate on every generation, in exchange for waiting until capacity frees up. The Google Cloud console bills in VM-hours, not chip-hours, so a multi-chip VM shows a bigger number than the table. Europe and Asia regions cost more: Ironwood is $13.20 on demand in London.
A chip-hour is not a unit of work though, and this is the part people skip. One Ironwood chip exposes two PyTorch devices, and a GPU, a Trainium2 chip and a TPU each finish a different amount of work per hour. The one third-party number Google cites is from SemiAnalysis's InferenceX benchmark, where TPU v7 with vLLM-TorchTPU served Qwen 3.5 at 2.01 cents per million tokens on the 8K/1K test. Treat that as a reason to benchmark your own model, not as your price.
To compare per-chip prices across the native PyTorch backends that now exist, pick an accelerator:
Pick a chip to compare list prices and backends. Prices come from the Google Cloud TPU and AWS Capacity Blocks pricing pages; a chip-hour buys a different amount of work on each.
For the GPU side of the comparison, my breakdowns of CoreWeave pricing and Lambda pricing cover the neoclouds, and the AWS pricing guide covers Trainium's home.
Benchmarks and model coverage
Google published more numbers with this launch than AWS did with TorchNeuron, and most of them is actually useful. Start with coverage, because for a 0.2.0 release "does my model run at all" matters more than speed. Per the PyTorch blog, Google continuously tests about 5,000 popular Hugging Face models in eager and compiled modes:
| Test | Result |
|---|---|
| Models tested end to end | 4,917 across 500+ architecture families |
| Run without crashing | 98% (4,821 of 4,917) |
| Numbers match GPU | 97% (638 of 657 representative models) |
| Bounded dynamism, correct numerics | 99% (834 of 839 causal LLMs) |
The numerical parity row is the one I'd look at first. At eesel, I learned the hard way that "it runs" and "it gives the right answer" are different tests: a support agent can look great in a demo and still give confident wrong answers on real tickets, which is why every rollout gets simulated against a customer's historical tickets first. Google checking output against a GPU baseline, not just checking for crashes, is the compiler version of the same habit.
On speed, three results stand out:
- Plain PyTorch, no custom kernels. A pure-PyTorch Muse Glimmer 30B written with standard modules reached 59% MFU on prefill on one TPU v7 chip and 38% MFU training on four. A version written for GPUs by the TorchTitan team, minimally ported, hit 36%.
- Close to JAX. Qwen3 MoE 235B in TorchTitan ran at 1,828 tokens per second on 64 TPU v7 chips, against 1,937 for the same model in JAX on MaxText. That's about 94% of Google's home framework.
- Scale. The same model scaled from 1,024 to 8,192 Ironwood chips with MFU moving only from 28.8% to 27.1%. Google's Cloud blog also cites up to 52% MFU on Qwen 3 8B.
Recommendation models got their own result: a DLRM v2 on 32 TPU v7x chips matched Google's tuned JAX version at an 18.3 ms step time, using TPU SparseCores for the embedding lookups. All of these are Google's own numbers on Google's own hardware, so I'd treat them as an upper bound until independent results show up.
What TorchTPU doesn't do yet
Google is direct about the gaps. The feature table splits everything into "Validated" and "Experimental," and the PyTorch blog lists the rough edges in detail.
- Dynamic shapes under compile.
torch.compileruns withdynamic=Falsetoday. Variable batch sizes or sequence lengths mean padding, bucketing or recompiles until bounded dynamism ships. - FlexAttention. Not supported yet; Google swapped it for
F.scaled_dot_product_attentionwhen porting a GPU model. - Distributed edge cases. FSDPv2 and code that differs by rank (MPMD) work but are marked "low perf." The default "skinny collectives" can cost 30-40% end to end; you win that back with a
@spmd_safedecorator or SimpleFSDP. - One process per TPU device. A TPU can only be opened by one process, so RL setups that colocate a trainer and a sampler on the same chips don't work out of the box.
- GPU-shaped models. TPU matrix units like dimensions of 128 or 256. Models that hard-code attention heads to 64, which is common, run but leave speed on the table.
- Model families still being tuned. GLM, Kimi, DeepSeek, GPT OSS, Mamba and LSTM/GRU are listed as experimental. Pre-training at full scale is "future exploration"; inference and fine-tuning are the production-ready jobs.
The bigger limit is not in the code at all. TPUs are Google Cloud chips (plus on-prem deals for very large customers), and the Hacker News launch thread, with 379 points and 93 comments, kept circling back to that:
"TPUs are unobtanium so while I appreciate Google's effort here I also recognize that it is not something that will be useful to me."
Another commenter, from a team that evaluated on-prem TPUs, put the buyer's view in one line:
"We've evaluated them but decided to stick with Nvidia, at least for one more generation."
And on the April preview, one reply framed the pitch more precisely than the marketing did: "just change one line and it works" was never the same promise as "just change one line and you will get peak performance on the TPU" (xrikcus, Hacker News). That matches Google's own docs, which is a good sign.
How TorchTPU compares to other ways to run models
For a team training or serving its own PyTorch models, the options now look like this:
| Option | What you manage | Code changes | Best for |
|---|---|---|---|
| TorchTPU on Cloud TPUs | TPU VMs or GKE, compile cache, shapes | "tpu", tpu_dist, static shapes for compile | Google Cloud teams with large, steady jobs |
| TorchNeuron on AWS Trainium | Instances, compile cache, shapes | "neuron", static shapes for compile | AWS-committed teams |
| NVIDIA GPUs (CoreWeave, hyperscalers) | Instances or clusters | None for CUDA code | Custom CUDA kernels, fastest start |
| Serverless GPUs (Modal, Baseten) | Your model code | Small wrappers | Bursty inference without capacity planning |
| Hosted model APIs (Groq, Fireworks AI) | Nothing | API calls only | Open models without touching hardware |
The first two rows used to say "rewrite in JAX" and "rewrite your training stack." Now they look a lot like the GPU rows, which is the whole point of both launches. In terms of choosing between them, it mostly comes down to which cloud you already pay, and also which chip is free when you need it.
If your end goal is an agent on Google Cloud rather than your own serving stack, Google's Agent Development Kit and the Gemini API sit a layer above all of this, and Gemini itself is trained and served on TPUs.
Who TorchTPU is for
TorchTPU is for teams that already write PyTorch, can get TPU capacity on Google Cloud, and spend enough on accelerators that a pilot pays for itself. Research teams get the most from eager mode and the Debug setting, since TPUs finally behave like any other device in a debugger. Inference teams get the most from vLLM-TorchTPU and SGLang-TorchTPU, which keep the same servers and APIs while swapping the hardware underneath.
It's not for you if your code depends on custom CUDA kernels, if you need dynamic shapes in compiled mode today, or if you can't get TPU quota. And it's not for you if you never touch model weights at all. Most of the companies I talk to at eesel are in that last group: they want AI to do a job, and which chip it runs on is someone else's decision.
Where eesel fits: this is infrastructure, eesel is the employee
TorchTPU makes the hardware layer cheaper to try and easier to switch. That's good for everyone who builds models, and over time it lowers what AI products cost to run. But it doesn't answer a ticket or write a blog post for you. Those are jobs, and jobs need context: your help center, your past tickets, your refund rules.
That's the layer eesel works at. eesel is an AI teammate platform: you hire ready-to-work teammates for specific jobs, today an AI helpdesk teammate and an AI blog writer. The helpdesk teammate joins your existing queue, learns from your help center and past tickets, and gets tested on your historical conversations before it answers a real customer.
SamCart shows what that looks like in practice. It runs eesel on Freshdesk and website chat, with a separate agent per channel answering from about 280 help center articles. The Freshdesk agent sends the first reply only until a person picks the ticket up, and a plan of 8,000 interactions a month covers it.
OroSoft, which makes trading software for bullion businesses, uses it the careful way: support staff mention eesel on a Freshdesk ticket and get a draft grounded in about 6,100 resolved tickets, with the source named. When nothing supports an answer, it leaves a note asking for a person instead of guessing. That's the same discipline as Google's numerical parity checks, applied to customer replies.
If you're the kind of reader who just installed a backend from a pinned pip index, you'll probably want to set this up from a terminal too. The eesel CLI operates the same teammate and workspace as the dashboard. npx @eesel/cli init chat-bubble --site https://your-site.com creates an agent, adds a website chat bubble and starts reading your site with no account, and eesel login keeps the workspace afterwards. From there, eesel integrations connect freshdesk wires up the helpdesk, eesel files upload adds knowledge, eesel chat talks to the agent and eesel approvals lets you approve held actions.
That design is why it works with coding agents, the same way Google's TPU agent suite plugs into them. An agent like Claude Code or Cursor can run the setup, read the results and fix what failed. Every workspace is also an MCP server, so the same agent can query eesel directly; see my Claude Code CLI reference if you're wiring that up.
The log above is what eesel activity prints in the terminal, newest first. Pricing is in units support teams already count: on the pricing page, a ticket or a chat is 1 credit, plans start at $299 a month for 500 credits, and the free tier gives you 100 credits to try everything.
Try eesel
TorchTPU lets you pick the cheapest chip for your own models. eesel lets you skip that decision for support work: a helpdesk teammate that plugs into Freshdesk, Zendesk or Gorgias, learns from your past tickets, and proves itself on historical conversations before it touches a live one. For SaaS teams with technical questions, the AI helpdesk API and the CLI mean you can script the whole setup. Try eesel free, or start from the terminal with npx @eesel/cli.
Frequently Asked Questions
What is TorchTPU?
tpu device to standard PyTorch, so eager mode, torch.compile and torch.distributed run on TPU v5e, v5p, Trillium and Ironwood. It was open-sourced on 20 October 2026. It is the chip layer; if you want AI that does a job, an AI teammate sits at the other end of the stack.Is TorchTPU free?
torch_tpu. You pay for the TPU time it runs on, from $1.20 per chip-hour for TPU v5e to $12.00 for Ironwood on demand in us-central1. My Gemini pricing guide covers the model side of Google's bill if you don't want to run chips at all.Does TorchTPU replace PyTorch/XLA?
torch.compile backends use. AWS made the same move for Trainium with TorchNeuron.Does TorchTPU really need only one line of code changed?
"cuda" to "tpu" and backend="nccl" to backend="tpu_dist". To run fast, no. Google's own docs say a model's shapes, sharding and ops may be tuned for GPUs, so peak speed takes compile work, wider attention heads or Pallas kernels. Coding agents like Claude Code can do much of that port, and Google ships a TPU agent suite for it.Which TPU generations does TorchTPU support?
Does TorchTPU work with Hugging Face Transformers?
device_map="tpu", and Diffusers runs with pipe.to("tpu"). Google reports 98% of 4,917 tested Hugging Face models run without crashing, and Hugging Face says Optimum TPU is no longer needed. More on the platform in my Hugging Face overview.How does TorchTPU compare to TorchNeuron on AWS Trainium?
torch.compile for speed, static shapes for now. TorchTPU covers four TPU generations and lowers to XLA; TorchNeuron covers Trainium2 and Trainium3 and lowers to the Neuron compiler. A Trainium2 chip-hour is $2.235 in EC2 Capacity Blocks, between Trillium and v5p on demand. See my AWS pricing breakdown for the rest of that bill.