BREAKING
AI Tools

TorchTPU vs TorchNeuron: The One-Line Escape From NVIDIA's CUDA

Updated October 26, 2026
12 min read
Editor's Note: This article was human-reviewed for accuracy
Blue-lit server blades in a data center rack

Quick Answer & Key Takeaways

  • • Google's TorchTPU and Amazon's TorchNeuron both launched at PyTorch Conference 2026, a day apart, and both promise the same thing: run your existing PyTorch code on a non-NVIDIA chip with a one-line device change.
  • • Both tools are free and open source, but Google lists TPU prices from $1.20 to $12.00 per chip-hour while Trainium2 has no on-demand price and sells only via prepaid Capacity Blocks at $2.24 per chip-hour.
  • • Neither tool ends chip lock-in yet, because the speed still comes from hand-tuned kernels written for one vendor's silicon.

On October 19, 2026, Amazon announced TorchNeuron, a native PyTorch backend for its Trainium chips. The next day, at the same PyTorch Conference in San Jose, Google announced TorchTPU, a native PyTorch backend for its TPUs. Within 24 hours, the two largest cloud vendors had shipped the same idea: your existing PyTorch code, running on their chips, with a one-line change and no rewrite. TechCrunch framed it bluntly — Google had "joined Amazon in attacking Nvidia's stronghold with AI app developers." The software is free on both sides. The hardware bills are not alike, and the fine print is where the two releases actually differ.

In short: TorchNeuron turns Trainium2 and Trainium3 into ordinary PyTorch devices with model.to('neuron'), in public beta since October 20, 2026, under Apache 2.0. TorchTPU does the same for Cloud TPUs from v5e to Ironwood with model.to('tpu'), open source since October 20, 2026, also under Apache 2.0. Google publishes per-chip-hour TPU prices from $1.20 to $12.00, while Trainium2 has no on-demand list price and sells through prepaid Capacity Blocks at $2.24 per chip-hour. Neither tool ends chip lock-in yet, because the speed still comes from hand-tuned kernels written for one vendor's silicon.

What each one actually is

Both tools plug into the same PyTorch extension point, PrivateUse1, the framework's official mechanism for registering custom hardware as a native device. That is why the pitch sounds identical: the chip shows up as a device string, standard PyTorch APIs work as written, and code written against the device-agnostic torch.accelerator API runs unmodified on either one.

TorchNeuron is Amazon's native PyTorch backend for AWS Trainium, announced on the PyTorch blog on October 19, 2026, and released as open source under Apache 2.0 at github.com/aws-neuron/torch-neuronx, going live as a public beta on October 20. It supports Trainium2 and Trainium3 instances only — first-generation Trainium and Inferentia are not covered. It offers eager execution for research and debugging, including an experimental "adaptive eager" mode that groups operations into sub-graphs to approach compiled performance without a code change, plus torch.compile for production, the Neuron Kernel Interface (NKI) for custom kernels, and Neuron Explorer profiling through the standard torch.profiler API.

TorchTPU is Google's native PyTorch stack for Cloud TPUs, announced on the Google Cloud blog on October 20, 2026, by Mark Lohmeyer and Bill Jia, and open sourced on GitHub under the google-pytorch organization, also under Apache 2.0. It was co-developed with the PyTorch team at Meta. It covers TPU generations v5e, v5p, v6e (Trillium), and v7 (Ironwood), supports eager debugging and torch.compile for production, and ships with native backends for vLLM and SGLang, the llm-d 1.0 distributed inference stack, TorchTitan for large-scale training, VeRL for reinforcement learning, Ray for orchestration, and LiteRT for on-device inference. It follows another recent Google Cloud AI release, the Data Agent Kit, which brings free MCP tools to more than 15 data services. In Hugging Face Transformers v5, deploying a model to a TPU is as simple as passing device_map="tpu". Google also shipped a TPU Agent Suite in early preview that automates onboarding, profiling, and kernel tuning through coding-agent harnesses including Claude Code and Antigravity.

The one-line switch, and where it stops being one line

The headline code sample in Amazon's announcement is three lines long, and the company calls the device change "the only change":

model = MyModel().to('neuron') # the only change

Google's version is the same shape: model.to("tpu"). In eager mode, both claims hold up on the vendors' own terms — operations dispatch individually, Python control flow and breakpoints work, and no compile or trace step stands in the way. That is the mode for research and debugging, and it is where the one-line promise is genuinely true.

Production is a different story, and both vendors say so in their own documentation. A compiled model on Trainium needs the Neuron backend named once, with torch.compile(model, backend="neuron") or torch.compiler.set_default_backend("neuron"). On TPU, production performance comes through torch.compile and custom Pallas kernels. The compile step is where the engineering moves: shapes, fusion, and kernel selection all live there, and a model that runs in eager mode is not the same as a model that runs fast.

TorchNeuron carries one more explicit limitation in its first release: native dynamic-shape support is still in development, so compiled models currently expect static shapes — inputs that change size mean recompiling. FP8 support is in private beta and reaches general availability in a later release. On the TPU side, Google's own launch material is quieter about gaps, but the same physics applies: out-of-the-box eager performance is far from tuned performance. Launch-week analysis of TorchTPU noted an untuned fine-tuning step taking 11.49 seconds against 2.78 seconds fully tuned — a fourfold gap between "it runs" and "it runs well."

What a chip-hour really costs

This is where the two releases stop looking like twins. Google publishes a full TPU price list; Amazon does not publish an on-demand price for Trainium2 at all.

Google's published per-chip-hour rates, from the Cloud TPU pricing page:

TPU generationOn-demandFlex-start3-year commitment
v5e$1.20$0.60$0.54
Trillium (v6e)$2.70$1.35$1.22
v5p$4.20$2.10$1.89
Ironwood (v7)$12.00$6.00$5.40

The fine print matters. Those are baseline US rates: Trillium costs $2.97 in Amsterdam and $3.24 in Tokyo, and Ironwood costs $13.20 in London. Flex-start, Google's Dynamic Workload Scheduler option, halves the on-demand rate but runs for windows of 10 minutes to 7 days and bills until deleted. A three-year commitment cuts Ironwood to $5.40, less than half the on-demand rate — the purchase model matters more than the generation. And the per-chip-hour number is not a per-token number: at the on-demand rate, a v6e serving one stream costs far more per million tokens than the same chip serving 32 concurrent streams, because utilization, not the sticker price, sets the token cost.

Trainium2 has no equivalent list. AWS sells Trn2 capacity through EC2 Capacity Blocks, prepaid reservations for a fixed window, at $2.235 per chip-hour — $35.76 an hour for a trn2.48xlarge with 16 chips, in US East (Ohio). Operating-system charges are billed separately, and reserving a window you do not fully use still bills the full window: you are buying a reservation, not consumption. There is no published per-hour rate for Trn3 UltraServers. For comparison on the same Capacity Block basis, an H100 costs about $5.97 per accelerator-hour, which puts a Trainium2 chip-hour at roughly 37% of an H100 hour. AWS's own claim on the Trn2 page is narrower and more honest than the raw price gap: Trn2 instances offer "30-40% better price performance than GPU-based EC2 P5e and P5en instances." A cheaper chip-hour is not a cheaper job — a Trainium2 chip and an H100 do different amounts of work per hour, and no independent head-to-head benchmark of the two new backends exists yet.

TorchTPU vs TorchNeuron, side by side

TorchTPU (Google)TorchNeuron (AWS)
AnnouncedOctober 20, 2026, PyTorch ConferenceOctober 19, 2026, PyTorch blog
StatusOpen source (v0.2.0 on PyPI)Public beta since October 20, 2026
LicenseApache 2.0Apache 2.0
HardwareCloud TPU v5e, v5p, v6e, v7 (Ironwood)Trainium2, Trainium3 only
The one linemodel.to("tpu")model.to("neuron")
ServingvLLM, SGLang, llm-d 1.0vLLM Neuron plugins, vLLM Omni Neuron
TrainingTorchTitan, VeRL, RayTorchTitan (torchtitan-neuron), TRL
Custom kernelsPallasNeuron Kernel Interface (NKI)
API accessNative PyTorch backend; no hosted API — you rent the chipsNative PyTorch backend; no hosted API — you rent the chips
Published chip priceYes — $1.20 to $12.00 per chip-hourNo on-demand price — Capacity Blocks at $2.235 per chip-hour
Installpip install torch_tpupip install torch-neuronx (AWS Neuron pip repo)

The evidence so far

Both vendors lead with numbers, and both sets of numbers are their own. Google's one training figure is that TorchTPU's training paths through TorchTitan reached up to 52% Model Flops Utilization on Qwen 3 8B, with near-linear scaling across up to 8,192 Ironwood TPUs. A customer quote in the same post — Zenteiq, running sovereign scientific foundation models on GKE — reports over 4x higher aggregate inference throughput using llm-d's prefix-cache-aware routing, from 1.47 billion to 5.97 billion tokens a day per 1,000 TPU v6e accelerators. A third-party data point, SemiAnalysis's InferenceX benchmark cited in launch-week coverage, put TPU v7 serving with the vLLM-TorchTPU stack at 2.01 cents per million tokens for specific inference tasks. None of these has an independent replication yet.

Amazon's one customer number is from Deepgram, the speech-recognition company: after a multi-month engagement with the AWS Generative AI Innovation Center, Trainium processed about 60% more input audio seconds per dollar than the company's existing GPU baselines. Cartesia, the voice AI company behind Sonic and Ink, reports running its standard inference pipeline with minimal changes to its PyTorch code, then writing NKI kernels for the layers that run most often and set end-to-end latency. Amazon also says it continuously runs 100,000 Hugging Face models — covering 99.9% of real-world usage by download count — against TorchNeuron, and publishes the gaps: its August 2026 upstream test run produced about 154,000 outcomes, with roughly 39,000 deliberate Neuron skips, the largest group being about 22,000 operators without a native Neuron lowering yet. That candor is the most useful number in either launch — it tells you exactly how unfinished the beta is.

The catches

Beta means beta, on the AWS side at least. TorchNeuron is a public beta, and its own documentation lists what the first release leaves out: dynamic and symbolic shapes, sparse tensors, NestedTensor, forward-mode autograd, and per-channel quantization, among others. Google does not label TorchTPU a beta, but version 0.2.0 on PyPI is a young release, and its TPU Agent Suite is explicitly in early preview.

The lock-in moved, it did not end. The code moves with one line; the speed does not. Cartesia's own account makes the split plain: small, fast ops map efficiently through the native path, but the layers that matter get hand-written NKI kernels. Those kernels run on Trainium and nowhere else. The same is true of Pallas kernels on TPU. A team that tunes hard on either chip has built an asset it cannot carry to the other — or back to NVIDIA — without rewriting the hot path. Before pricing a move, ask which of your workloads would need hand-tuned kernels, who writes them, and what moving back would cost.

There is no head-to-head yet. No publication has run the same model through TorchTPU and TorchNeuron on the same workload and compared throughput, latency, and invoice. The per-chip-hour arithmetic in this article is a starting point, not a verdict: $2.235 for a Trainium2 chip-hour against $2.70 for a Trillium chip-hour tells you nothing about which one serves your model cheaper, because the chips differ in memory, compute, and the maturity of the software around them.

Both still need the cloud. Neither tool is something you install on hardware you own. TorchTPU runs on Google Cloud TPUs; TorchNeuron runs on EC2 Trn2 and Trn3 instances. The "open source" label covers the software stack, not the silicon — and on the AWS side, access itself is rationed through prepaid Capacity Blocks rather than on-demand rental. The Trainium push sits alongside AWS's other infrastructure bets, including AWS Infinity, the company's open-source runtime for running tens of thousands of agents.

Pros and cons

TorchTPU pros: published, transparent per-chip-hour pricing with steep commitment discounts; broadest serving stack (vLLM, SGLang, llm-d 1.0); Hugging Face Transformers v5 support with device_map="tpu"; training scale demonstrated to 8,192 Ironwood chips; Meta already runs production workloads on it. Cons: young v0.2.0 release; Ironwood at $12.00 on-demand is the most expensive chip-hour on either side; TPU Agent Suite still in early preview; no independent head-to-head benchmark.

TorchNeuron pros: lowest per-chip-hour price of the two ($2.235 vs $2.70 for the comparable Trillium); Deepgram's 60%-per-dollar figure is a real customer result, not a synthetic benchmark; unusually transparent test-gap reporting (100k-model sweep, published skip counts); NKI gives fine-grained kernel control. Cons: public beta with real feature gaps (no dynamic shapes, FP8 still private beta); no on-demand pricing — Capacity Blocks only, prepaid and billed whether used or not; covers only Trainium2/3, leaving Trn1 and Inferentia behind; no independent head-to-head benchmark.

Who should try which

If your workloads already live on Google Cloud, TorchTPU is the lower-friction experiment: install the package, point an existing PyTorch script at a v5e or v6e, and measure. The $1.20 v5e on-demand rate — $0.60 on flex-start — is the cheapest on-ramp to either chip family, and it fits models up to roughly 27B parameters. If you live on AWS, TorchNeuron is the equivalent test, but budget it as a reservation, not a rental: Capacity Blocks are prepaid windows, so a trial needs a planned window and a workload ready to fill it. Teams with no cloud commitment should try the cheaper chip-hour first and let the invoice decide — but only after porting one real workload end to end, kernels included, because eager-mode demos flatter both tools equally. And teams whose models depend on dynamic shapes should wait on TorchNeuron until that support lands.

FAQ

Do TorchTPU and TorchNeuron really need only a one-line change?
In eager mode, yes, on the vendors' own terms: model.to("tpu") or model.to("neuron") moves the model, and standard PyTorch APIs work as written. Production performance is a separate job — the compile step, kernel tuning, and shape handling are where the real work lives, and TorchNeuron's first release does not support dynamic shapes at all.

Which one is cheaper?
Per chip-hour, Trainium2 at $2.235 beats Trillium at $2.70 on-demand, but Google's v5e undercuts both at $1.20, and Ironwood at $12.00 costs more than five Trainium2 chip-hours. Per job, nobody knows yet: no independent benchmark has run the same workload on both, and throughput per dollar depends on the model, the kernels, and how well the chip is utilized.

Can I run the same model on both?
The model code, yes — code written against PyTorch's device-agnostic torch.accelerator API runs unmodified on either backend, and both support Hugging Face Transformers v5. The performance-critical kernels, no — Pallas kernels target TPUs and NKI kernels target Trainium, and neither ports to the other.

Is this the end of CUDA lock-in?
Not yet. Both tools remove the porting tax on model code, which is real progress, but they move the lock-in into the kernel layer: the hand-tuned code that makes a workload fast is written for one vendor's silicon. The honest test of portability is not the first run, it is the second move.

Which should I try first?
The one that matches the cloud you already pay for. TorchTPU if you are on Google Cloud, TorchNeuron if you are on AWS. Both are free, open-source software — the cost is chip time, and on the AWS side that means a prepaid Capacity Block window, so plan the trial before you book it.

Key takeaways

• Google's TorchTPU and Amazon's TorchNeuron launched a day apart at PyTorch Conference 2026, both turning a rival chip into a native PyTorch device with a one-line change, both Apache 2.0 open source.
• The one-line promise is real in eager mode and incomplete in production: compile configuration, kernel tuning, and — on Trainium — static shapes are where the work moves.
• Google publishes TPU prices from $1.20 to $12.00 per chip-hour with commitment discounts up to 55%; Trainium2 has no on-demand price and sells through prepaid Capacity Blocks at $2.235 per chip-hour.
• The only real customer numbers so far are Deepgram's roughly 60% more audio per dollar on Trainium and Zenteiq's 4x throughput gain on TPUs — both vendor-selected, neither independently replicated.
• Chip lock-in did not end; it moved into the kernel layer, so the deciding test is a real workload ported end to end, measured on the invoice, not the announcement.

Both tools are live now — TorchNeuron as a public beta, TorchTPU as open source on GitHub — and the honest test is the same for both: port one real workload, measure it, and compare the invoice, not the announcement. As of October 26, 2026, nobody has published that head-to-head yet.

12 min read
15 views
edit_note Written by

Ratul Ahamed

Founder & Editor-in-Chief

Ratul Ahamed is the Founder and Editor-in-Chief of AI Core Updates, steering the publication's strategic growth, technological architecture, and human-verified journalistic standards. With over eight years of hands-on experience across full-stack systems and machine learning infrastructure, Ratul founded AI Core Updates to deliver transparent, benchmark-backed intelligence that bridges the divide between foundational research and real-world executive decision-making. At AI Core Updates, he directs long-term editorial investigations and leads the development of internal automated verification workflows.

Read all stories by Ratul Ahamed arrow_forward
verified_user Human-Reviewed Editorial

forum Comments (0)

Leave a Reply

chat_bubble_outline

No comments yet. Be the first to join the conversation!

AI Core Updates Insider Dispatch

Stay Ahead of the Exponential AI Curve

Join over 48,000+ AI researchers, software engineers, and founders. Get our daily 5-minute morning intelligence briefing analyzing breakthrough research, frontier models, and autonomous workflows.

bolt
Daily Briefing in 5 Mins
High signal, zero marketing fluff.
podcasts
Podcast Deep Dives
Early access & episode transcripts.
psychology
Evaluated Prompts & Tools
Tested benchmarks & system designs.
verified_user
Strict Anti-Spam Policy
1-click unsubscribe at any time.
AK
SL
MR
JD
Trusted by builders from leading AI labs & enterprises.
mark_email_unread

Subscribe to Newsletter

Delivered daily at 7:00 AM UTC

person_outline
mail_outline
Include in my subscription:

No spam ever. You can easily unsubscribe with a single click at any time.

cookie

Cookie Settings

We use cookies to enhance your browsing experience, serve personalized content, and analyze our traffic. By clicking "Accept All", you consent to our use of cookies. Learn more