News · · 9 min read

Google Open-Sources TorchTPU: The Defaults That Change When PyTorch Runs on TPUs, and Why Tokyo v6e Matters

In this post
  1. What changed since April
  2. Two backends in the same week
  3. Between "it runs" and "it runs fast"
  4. One third-party measurement, for serving
  5. The nearest TPU to Seoul
  6. What changes for teams in Korea
  7. Community reaction
  8. Limitations
  9. Sources

Google open-sourced TorchTPU at PyTorch Conference 2026 on October 20, 2026. TorchTPU registers Google's TPUs as a PyTorch device (tpu), and Google built it together with the PyTorch team at Meta. It is licensed under Apache 2.0, and version 0.2.0 is on PyPI. Porting starts with changing model.to("cuda") to model.to("tpu"). Two things matter before that change, though. First, TorchTPU's default matmul precision uses BF16 inputs, so code that ran in FP32 on CUDA can produce slightly different numbers. Second, Google Cloud's TPU region list has no TPUs in Seoul, and the only Asia-Pacific zone offering a generation TorchTPU validates (v5e and later) is v6e in Tokyo.

AWS announced TorchNeuron for Trainium one day earlier. This post covers what the open-source release adds, how the two backends compare on the same axes, the distance between "it runs" and "it runs fast," the one third-party measurement, and the region picture for teams in Korea.

What changed since April

TorchTPU is not a new name. Google described its design on the Google Developers Blog on April 7, 2026, and that post listed a public GitHub repository as work for later in 2026. The October 20 release is the point where you can actually install it. The google-pytorch/torch_tpu repository is public, pip install torch_tpu==0.2.0 installs it, and Python 3.11 or later is supported.

The release covers more than one backend. The PyTorch blog says four repositories open at this conference: TorchTPU, vLLM-TorchTPU, SGLang-TorchTPU and TorchTitan. Google Cloud's post adds Ray support, a TPU backend for the VeRL reinforcement-learning framework (in preview), and a TPU agent suite that automates onboarding, diagnostics and inference performance tuning (in early preview). Hugging Face says Transformers v5 can place a model on a TPU with device_map="tpu", and that Optimum TPU is no longer required.

The previous path explains why this matters. Since 2020, the official way to run PyTorch on TPUs was PyTorch/XLA. The PyTorch blog says plainly that it required intrusive source changes to mark compilable regions, suffered from long compile times and did not support eager execution. TorchTPU was built the other way around: eager execution first, and no JAX dependency in the runtime.

Two backends in the same week

AWS announced TorchNeuron on October 19 (its open-source public beta opened October 20), and Google released TorchTPU on October 20. Both register the vendor's own chip through PyTorch's official extension point, PrivateUse1. We covered TorchNeuron's code changes and constraints in our TorchNeuron post. We compiled the table below from each vendor's announcement and README.

ItemGoogle TorchTPUAWS TorchNeuron
Release formOpen source, PyPI 0.2.0Open source, public beta
LicenseApache 2.0Apache 2.0
Device nametpuneuron
Target hardwareValidated on TPU v5e and laterTrn2, Trn3
Python3.11 and later3.12 to 3.14
Hugging Face sweepAbout 5,000 modelsAbout 100,000 models
torch.compile shapesCurrently forced to dynamic=Falsedynamic=False only

The two start from the same place, and their first releases share a constraint. Both begin with static shapes, recompiling whenever the input shape changes. TorchTPU calls this a temporary limitation and is building bounded dynamism, where the maximum size is fixed ahead of time. According to the PyTorch blog, a test on 839 causal language models produced correct numerics on 834 (99%).

The sweep sizes are not comparable as they stand. TorchNeuron's roughly 100,000 models are chosen to cover real-world usage by download count. TorchTPU's roughly 5,000 are drawn from the most-downloaded models, but Google published pass rates for them. Of 4,917 models, 4,821 (98%) ran without crashing, and 638 of 657 architecture representatives (97%) matched a GPU baseline numerically. The two vendors measured different things, so neither number shows which backend covers more.

Between "it runs" and "it runs fast"

The PyTorch blog states up front that changing to model.to("tpu") will not reach peak TPU performance, because a model's shapes, sharding and ops may be tuned for GPUs. The one-line change promises execution, not speed. Google's FLUX training case shows how large that gap is. Figure 1 compiles the step time as each optimization is added.

Figure 1. TorchTitan FLUX training step time, from 11.49 s with default eager to 2.78 s with all optimizations. Google-reported figures from the PyTorch blog, compiled by Kortress

The largest single-step multiplier came from torch.compile. The speedup over default eager went from 2.07x to 3.80x once compilation was on, taking that step from 5.55 s to 3.02 s. In absolute seconds, the SimpleFSDP sharding step saved the most, from 9.95 s to 6.52 s. With the final model-specific graph rewrites, 11.49 s dropped to 2.78 s; we computed that 4.13x by dividing the two values. The Pallas SplashAttention step in the middle is a TPU-specific kernel. The final number in this case reflects TPU kernels and compiler settings, not a one-line change.

Google also reported numbers for less-tuned code. A Muse Glimmer 30B model written only with standard PyTorch modules trained on four TPU v7 chips at 38% MFU (model FLOPs utilization) with a sequence length of 5,120. A TorchTitan implementation originally written for GPUs, with FlexAttention swapped for standard attention and torch.compile and SimpleFSDP applied, reached 36%. The blog notes that FlexAttention is not yet supported in TorchTPU. All of these figures are Google's own measurements.

The precision default looks like the difference most likely to be missed when porting. In the PyTorch blog's comparison table, CUDA's default precision mode is highest (pure FP32), while TorchTPU defaults to medium. Medium converts multiplication inputs to BF16 and accumulates in FP32. You can choose high or highest for near-FP32 results, at roughly 3x and 6x the compute cost. If you move a training script whose loss curve you already matched on GPUs, this is the first thing to suspect when the curves diverge.

We found no third-party measurement of TorchTPU training performance. The third-party data that exists is for inference serving. SemiAnalysis, a semiconductor research firm, added TPU v7 (Ironwood) to its InferenceX benchmark as a preview on September 7, 2026. The workload was Qwen3.5 397B in FP8, served by the TorchTPU-based vLLM stack, which was in private beta at the time.

The results cut both ways. SemiAnalysis summarized that TPU v7 delivered up to 50% better performance per dollar than B200 and B300 in FP8 aggregated serving. At the specific data point of 20 tokens per second per user, TPU produced 50.4% more tokens per dollar than B200 and 96.0% more than B300, and SemiAnalysis notes that this 50% to 96% advantage applies to that data point rather than every latency target. In the same paragraph, SemiAnalysis warns that the high-concurrency results come with a latency cost: at 256 concurrent requests, mean time to first token was 5.41 s on TPU, 3.75 s on B200 and 2.40 s on B300, the slowest of the three. In FP4 serving, NVIDIA GPUs keep the lead because TPU v7 has no native FP4 compute. Comparing aggregated TPU v7 with disaggregated GB300 NVL72 using Google's internal TCO, SemiAnalysis put GB300 about 30% ahead on performance per dollar in the middle of the latency range.

This measurement is not fully independent. In the same post, SemiAnalysis thanks 15 named Google staff and the Inferact team and labels the results "Official Preview." The PyTorch blog calls it a "third-party inference audit." It reads best as a third-party measurement prepared together with the vendor.

The nearest TPU to Seoul

To use TorchTPU you have to go where the TPUs are. We compiled the zones offering generations TorchTPU validates (v5e and later) from Google Cloud's TPU regions page, last updated October 20, 2026.

AreaTPU7x (Ironwood)v6ev5e and v5p
North America2 zones in us-central15 zones5 zones
EuropeNoneeurope-west4europe-west4
Asia PacificNoneasia-northeast1 (Tokyo)None
South AmericaNonesouthamerica-west1None

In Asia Pacific, the only zone with a TorchTPU-validated generation is v6e in Tokyo. Seoul (asia-northeast3) has no TPUs at all. The only other Asia-Pacific TPU is v2-8 in asia-east1 in Taiwan, and v2 is outside the generations TorchTPU says it validates. TPU v7 (TPU7x in the docs), the chip behind most of Google's published figures, is available only in the central US. The page adds that configurations with more chips are available only in limited quantities. Check actual availability in the console before planning around it.

What changes for teams in Korea

TorchTPU lowers the cost of trying TPUs. Until now, a PyTorch team had to rework its code around PyTorch/XLA or learn JAX. Now it can run existing code eagerly and see where it breaks first. Teams built on Hugging Face models can start a first experiment with device_map="tpu".

The region picture runs opposite to TorchNeuron's. Trn2, which TorchNeuron requires, is in neither Seoul nor Tokyo. TorchTPU has nothing in Seoul either, but Tokyo has v6e. Services that must stay in the Seoul region cannot use either backend directly. For teams that can use an overseas region, Tokyo v6e appears to be the closest option of the two chips in terms of latency. Serving that handles personal data needs a cross-border transfer review first.

A testing order follows from this. Step one is matching accuracy against GPUs in eager mode, keeping the BF16 default in mind. Step two is choosing a sharding strategy and turning on torch.compile, while fixing input shapes with padding or buckets. In the FLUX case, sharding and compilation delivered most of the speedup, and the TPU-specific kernel contributed less (6.52 s to 5.55 s), so kernel work appears safe to defer. For text generation with variable lengths, a serving engine that Google and its partners are tuning for TPUs, such as vLLM-TorchTPU, looks more realistic than tuning by hand.

Community reaction

On Hacker News, the post linking the GitHub repository, submitted on release day, had 333 points and 73 comments when we checked. Supportive comments compared it with past experience. User ismailmaj wrote: "I remember the days when you had to fight XLA to make PyTorch run on TPUs, it seems we finally got first-party support!" Another user, mathisfun123, read TorchTPU as a replacement for PyTorch/XLA.

The cautious side asked where to get the hardware. One comment asked, "Are there TPUs outside Google Cloud?" and another, "where can I buy a TPU?" User atty, replying about Google's on-premises TPU sales and describing their company as a relatively large GCP user, wrote: "We've evaluated them but decided to stick with Nvidia, at least for one more generation." The reaction suggests that once the software barrier drops, hardware access becomes the next one. We found no substantive discussion on Reddit or X as of checking.

Limitations

ItemDetail
Hands-on runsWe did not run TorchTPU on TPUs. Tables and figure compile official figures
Performance numbersAll training numbers are Google's. The only third-party data is for serving, a preview prepared with the vendor
FLUX caseThe case as a whole was evaluated on TPU v6e and v7x, but the step-time measurements do not state their generation or chip count
Comparison tableThe vendors validate differently, so the Hugging Face sweep sizes are not on the same basis
RegionsBased on the page updated 2026-10-20. Actual capacity in the console may differ
PricingThis post does not cover TPU pricing, which varies by region

This comparison reflects the first public releases. TorchTPU lists production-ready bounded dynamism and FP8/FP4 support for next-generation TPUv8 as next steps, and TorchNeuron lists general availability of FP8. TorchNeuron says only that it deliberately left dynamic shapes out of its first release. The constraint rows in the table may change with each release.

Sources

1 / 1