Google Cloud’s blog, Bill Jia, vice president of engineering, dated the announcement to PyTorch Conference 2026: TorchTPU, the native PyTorch stack for Google TPUs, is open source on GitHub. TPUs join GPUs and CPUs as a native PyTorch option. That is a device string. It is not cuda.
We already filed TorchNeuron public beta that was not a CUDA port and PyTorch 2.14 zeroing a clamp gradient. This week the named package is torch-tpu 0.2.0. The named trap is pip pulling a CUDA wheel of torch if you skip the CPU index.
0.2.0 is October 19. The keynote is October 20.
Unite.AI on October 20: the torch_tpu repository is public under the google-pytorch GitHub organization, Apache 2.0. Python 3.11 and later. torch-tpu 0.2.0 hit PyPI on October 19, 2026, Linux x86-64 wheels for CPython 3.11 through 3.14. Docs at torchtpu.readthedocs.io. Interactive notebooks have to run from a TPU VM; a v6e GCE instance is the easiest option.
The package date is the day before the conference post. If your CI tried pip install torch-tpu on October 19, you were on the wheel before the blog. If your CI runs on macOS ARM in GitHub Actions, you do not have that Linux x86-64 wheel. You have a laptop. The README’s notebooks are not for that laptop.
MarkTechPost the same day prints the install:
pip install torch_tpu
import torch
device = torch.device("tpu")
x = torch.randn((4, 4), device=device)
Then the sentence that should be in the pin file: the README pins torch_tpu==0.2.0 with the CPU PyTorch wheel index. Without it, pip pulls the CUDA build of torch. CUDA torch on a TPU VM is how you spend an afternoon importing the wrong backend. Device "tpu" will not save you if torch is the NVIDIA build.
Python 3.11 through 3.14 wheels means 3.10 is not in that list. If the image is still 3.10 because a data team froze it in 2024, this stack is not your image. Do not --use-deprecated your way around a wheel that does not exist.
Apache 2.0 is the license you can actually use. It is not a support contract. Google Cloud’s list of “tools you already rely on” (vLLM, SGLang, llm-d, TorchTitan, Ray, TorchRec, Hugging Face, LiteRT) is a compatibility pitch. Pin the versions you tested. A keynote list is not a lockfile.
Eager first. PrivateUse1. Fused is the 50 to 100 percent line.
Unite.AI recaps the April 7, 2026 introduction: Eager First on PyTorch’s PrivateUse1 interface; Debug, Strict, and Fused eager modes; shared compilation cache. Google said Fused Eager delivers a 50 percent to 100+% performance increase over Strict Eager. torch.compile on TorchTPU captures FX graphs with Torch Dynamo and lowers them through XLA and StableHLO, with custom kernels in Pallas and JAX. Distributed Data Parallel, FSDPv2, and DTensor out of the box. Architected for divergent execution across ranks. A public GitHub repo was a 2026 roadmap item in April. October is that item landing.
MarkTechPost: TorchTPU replaces PyTorch/XLA with eager-first PrivateUse1. torch.compile(backend="tpu") lowers through StableHLO and XLA. Device string tpu versus XLA’s xla versus Neuron’s neuron. If your code still says xm.mark_step(), you are in the stack this is meant to replace. Do not mix xla device code and tpu device code and call it a port.
Debug versus Strict versus Fused is a mode switch, not a free lunch. Debug is how you see the op that died. Fused is the 50–100 percent claim over Strict. If you benchmark Debug and ship the number, you shipped the slow path. Name the mode in the result.
Custom ATen kernels, compilation cache for XLA computations, data and tensor parallelism, debug and benchmark utilities: README inventory. Examples listed: Llama, Qwen, ResNet, DLRM, minGPT, UNet. An example is not your tokenizer. If you train ads ranking, Jana van Greunen’s quote is the closer one.
Jia’s post: co-developed with the PyTorch team at Meta. Debug in eager, then torch.compile for TPU utilization and memory overhead. Standard APIs, diagnostics, native distributed primitives. The sell is “minimal code changes.” Minimal is not zero. device="tpu" is a change. The wheel index is a change. The VM is a change. CUDA graphs you copied from a blog will not map.
52 percent MFU on Qwen 3 8B. 8,192 Ironwood. Not your box.
MarkTechPost: TorchTitan hit 52 percent MFU on Qwen 3 8B, scaling to 8,192 Ironwood chips. That is a cluster number. It is not a Colab. If you have one v6e, you do not have 8,192 Ironwood. Do not put 52 percent MFU on a slide next to a single-chip laptop demo.
vLLM-TorchTPU and SGLang-TorchTPU show up in that comparison table as the serving story, against vLLM Neuron on AWS and vLLM-JAX / SGLang-JAX on TPU. Serving stacks are named forks. If you pip install vllm and expect TPU, you installed the default. Read which extra you got.
Their table, compressed: TorchTPU, Google plus Meta, Cloud TPU v5e and later tested, device tpu, PrivateUse1, eager default with Strict/Debug/DeferAndFuse, Dynamo to StableHLO to XLA, DDP/FSDP2/DTensor/SPMD and MPMD, Apache 2.0, v0.2.0. PyTorch/XLA: device xla, lazy tensors, SPMD only, being replaced by TorchTPU. TorchNeuron: Trainium 2 and 3, device neuron, public beta Neuron 2.33. JAX on TPU: not PyTorch.
That last row is the fork people still miss. If the team already writes JAX, TorchTPU is not a reason to rewrite. If the team has PyTorch and a TPU quota, this is the native path Google wants instead of XLA’s lazy tensors.
We already told you Neuron is not CUDA. TorchTPU is also not CUDA. Two vendors, two device strings, same lesson. nvidia-smi will not help. The TPU VM is the machine. Notebooks that are not on that VM are screenshots.
Meta’s Jana van Greunen, director of PyTorch engineering: they cannot be constrained by a single hardware architecture; TorchTPU lets them deploy PyTorch workloads to Google TPUs for LLM, ads, and recommendation production use cases and give feedback. That is Meta’s fleet. Your fleet is the quota project.
Pin the CPU torch. Do not import on CUDA by accident.
MarkTechPost’s warning is the operational one. pip install torch_tpu without the CPU PyTorch index pulls CUDA torch. On a TPU VM that may still “work” until the first .to("tpu") or until you OOM a GPU that is not there. Put the index URL from the README in the Dockerfile. Pin torch_tpu==0.2.0. Pin the torch CPU wheel they name. Do not let dependabot bump you onto a CUDA torch because the version number looked newer.
Install on the TPU VM, not on the builder that runs unit tests on CPU-only GitHub runners unless you have a CPU fallback path you actually test. PrivateUse1 backends that only exist on the vendor VM are how CI goes green on dummy tensors and red on the job.
If you still have PyTorch/XLA in the same image, pick a winner. Two stacks that both think they own compilation will give you the worst of lazy and eager. MarkTechPost’s status line is explicit: XLA is being replaced by TorchTPU. Replacement is a migration. Migrations need a freeze window.
Custom kernels in Pallas and JAX mean someone on the team has to read those files when an op is missing. Missing ops are the actual port tax. “Minimal code changes” assumes your graph is in the supported set. The example list is the supported-set advertisement. DLRM in the examples is the ads hint. UNet is the image hint. Your custom autograd Function is the thing that will fail first.
LiteRT in Jia’s tool list is a mobile/edge name sitting next to vLLM. Do not assume the same wheel. Different runtimes. Same keynote sentence.
TorchTitan in that same list is the MFU source for the 52 percent Qwen 3 8B figure. If you are not on Titan, you do not get to borrow 52 percent. Your trainer is your trainer. Record tokens per second on the VM you have.
April 7 already promised divergent execution across ranks. October’s open source is that design in a repo you can read. Divergent ranks are how a hang becomes “works on 1, dies on 8.” If you scale past one v6e, test a two-chip job before you quote Titan’s 8,192. FSDPv2 and DTensor in the README are the APIs. Your wrapper around accelerate may still assume CUDA device IDs. Device IDs are the port.
Ray and TorchRec in Jia’s list are the production names next to Llama examples. Recsys on TPU is the Meta quote’s ads line. If you do not run TorchRec, that bullet is not yours. Hugging Face through standard PyTorch is the one most readers will actually try: a from_pretrained on a TPU VM. Try a tiny model first. If the loader pins CUDA kernels, you are back to the wrong torch wheel.
CPython 3.14 on the wheel list is ahead of a lot of images. 3.11 is the floor. If you are on 3.12 because that is what the TPU image ships, stay there until you prove 3.14. Do not chase the upper bound as a flex.
DeferAndFuse in MarkTechPost’s eager-mode list is another name next to Debug and Strict. If the README treats DeferAndFuse as the fused path, name it in the benchmark. If you cannot find it, you are not on the path they measured the 50–100 percent against.
What to type Monday
If you do not have a TPU VM, stop. The GitHub repo is not a GPU speedup.
If you do, read the README pin, install 0.2.0 with the CPU torch index, torch.device("tpu"), run the smallest example they ship (minGPT before Llama) and record MFU on your chip count, not 8,192. Check whether torch.compile(backend="tpu") is required for the number you care about. Name Debug versus Fused in the log.
Keep Neuron and CUDA jobs on their pins. This is an extra device, not a replacement for the A100 you did not return. Transformers loading GGUF still leaves llama.cpp on the box with no Python. TorchTPU does not change that box.
Do not merge this pin with the October 21 LiteLLM incident on the other page. Different package, different index, same rule: read what pip actually fetched. pip debug / pip show torch / pip show torch_tpu on the VM, not on your laptop.
cuda is a string your code already has. tpu is a new one. 0.2.0 is a wheel. A conference is a date. Pin the wheel.
Discussion
Leave a comment
No comments yet
Be the first to start the conversation.