AIAIExplainer NO. ---

TorchTPU and TorchNeuron explained: PyTorch beyond Nvidia CUDA

Both tools plug into PrivateUse1, PyTorch's official extension point for new hardware, so moving a model can be as simple as changing the device setting. Neither is complete yet: Amazon lists operator gaps, including dynamic shapes.

By Rowan Ashby Published 21.10.26 12:19 UTC 9 min read

Close-up of a Google Tensor Processing Unit 3.0 board with its chips and heat sinks — Zinskauf / Wikimedia Commons / CC BY-SA 4.0 · Source

Key takeaways

Google has released TorchTPU as an open-source project on GitHub. It is a native PyTorch stack for its Tensor Processing Units (TPUs), co-developed with Meta's PyTorch team.
A day earlier, Amazon announced TorchNeuron, an Apache 2.0-licensed native PyTorch backend for its Trainium2 and Trainium3 chips. It entered public beta on 20 October.
Both tools plug into PrivateUse1, PyTorch's official extension point for new hardware. In practice, moving a model can be as simple as changing the device setting.
Neither tool is complete yet. Amazon lists operator gaps and features it does not support, including dynamic shapes. Google's own engineers have said changing input sizes triggers recompilation.
The race is about software as much as silicon. Nvidia's advantage rests heavily on CUDA, the programming platform that PyTorch was largely built around.

TorchTPU and TorchNeuron are native PyTorch backends that let PyTorch code run on Google's Tensor Processing Units and Amazon's Trainium chips with little or no rewriting. Amazon released TorchNeuron first; Google followed a day later with an open-source TorchTPU release.

This explainer covers what a PyTorch backend is, what each company released, what developers can and cannot move today, and why the software layer may decide more of the AI hardware race than chip specifications.

First, what is PyTorch and why does it matter?

PyTorch is an open-source framework for building and training neural networks. Meta (then Facebook) created it in 2016 and handed it to the Linux Foundation in 2022, according to TechCrunch. Most large language models, image generators and speech systems in research today are defined in PyTorch code. Amazon says 85% of AI researchers and developers write their code in PyTorch. That figure comes from Amazon's own announcement and is not an independent count.

The catch is that PyTorch grew up on Nvidia GPUs. Many of its features, tests and performance tricks assumed CUDA, Nvidia's programming platform. Nvidia describes CUDA as "the foundation for GPU computing", a toolkit of compilers, libraries and debugging tools. A team that wanted to run PyTorch on another chip usually had to use special wrappers, adapt its code, or accept slower performance. Nvidia's grip also shapes the GPU-only cloud providers covered in our explainer what is a neocloud.

What is a "backend", in plain terms?

A useful way to picture it: PyTorch is the car, and the backend is the engine underneath. Developers drive using familiar controls, such as layers, tensors, optimisers and training loops. The backend converts those instructions into work a specific chip can execute. Amazon uses a similar analogy, describing TorchNeuron as a way to swap the engine without "tearing down the car".

The technical hook is a mechanism called PrivateUse1. According to PyTorch's own documentation, it is a reserved slot that lets hardware makers register a new device without editing PyTorch's core files. A backend that uses it registers its own versions of PyTorch's basic operations. Operations it has not yet implemented can fall back to the CPU. The device can also be renamed, so a developer writes device="neuron" or a TPU device name instead of something generic.

Both new tools use this route. That matters because earlier efforts, such as Google's older PyTorch/XLA project, sat beside PyTorch rather than inside it. Google engineers described the goal of TorchTPU in an April 2026 developer blog post with a simple phrase: "it should feel like PyTorch".

What Amazon released: TorchNeuron

Amazon engineers explained TorchNeuron on the official PyTorch blog. In the simplest case, the change is one line: moving the model to the 'neuron' device. Code written against torch.accelerator, PyTorch's recommended device-neutral interface, runs unmodified. A compiled model only needs the Neuron compiler backend named once.

According to that post and the AWS Neuron documentation, TorchNeuron supports:

Eager mode, where operations run immediately, for research and debugging, plus an experimental "adaptive eager" mode that fuses operations to cut overhead;
torch.compile, which captures a model as a graph and passes it through Amazon's Neuron compiler for production performance;
standard distributed training tools, including DDP, FSDP2, tensor and pipeline parallelism, and DTensor;
a kernel language, NKI, for hand-tuning the hottest parts of a model, and an open-source library of ready-made kernels.

Hugging Face says Trainium is now a native backend for its Transformers and Diffusers libraries. Amazon also reports that it contributed pull requests to PyTorch itself to remove "hard coded CUDA assumptions". TechCrunch described this as making PyTorch more hardware-neutral by default.

What Google released: TorchTPU

Google announced the open-source release at PyTorch Conference 2026 in a Google Cloud blog post. It was written by Mark Lohmeyer, vice president and general manager of AI and computing infrastructure, and Bill Jia, vice president of engineering. The code is on GitHub under the google-pytorch organisation. Google says existing PyTorch workloads can run on TPUs "with minimal code changes, using the familiar APIs".

Like TorchNeuron, TorchTPU offers an eager mode for debugging and a torch.compile path for production. Google's April engineering post explained that its compiled path uses XLA, Google's long-standing TPU compiler, rather than PyTorch's default Inductor compiler. It also described three eager variants, including a "fused eager" mode that Google says ran 50% to more than 100% faster than its strict single-operation mode.

Google's announcement focuses on the wider ecosystem. It lists native support in the vLLM and SGLang serving engines, the TorchTitan training library and the Ray orchestration framework. In Hugging Face Transformers v5, targeting a TPU is a device_map="tpu" setting. Google also released a preview "TPU agent suite" that helps automate onboarding, diagnostics and tuning. It updated LiteRT, its on-device runtime, so that PyTorch models can be converted for phones and laptops.

Side by side

Item	Amazon TorchNeuron	Google TorchTPU
Target chips	Trainium2, Trainium3	Google TPUs (Ironwood/v7x cited for large-scale tests)
Status	Public beta from 20 Oct 2026	Open source on GitHub; agent suite in preview
Licence	Apache 2.0	Not stated in Google's announcement
Integration route	PyTorch PrivateUse1	PyTorch PrivateUse1
Compiler path	TorchDynamo → StableHLO → Neuron compiler	TorchDynamo → StableHLO → XLA
Headline vendor figure	~154,000 upstream test outcomes in Aug 2026 run; ~69,000 skipped, of which ~22,000 for operators not yet lowered	Up to 52% model FLOPs utilisation on Qwen 3 8B, near-linear scaling to 8,192 Ironwood TPUs

Sources: PyTorch blog (TorchNeuron, updated 20 Oct 2026); AWS Neuron documentation; Google Cloud blog (TorchTPU, Oct 2026); Google Developers Blog (Apr 2026). All performance and test figures are company-reported and have not been independently benchmarked.

What developers can move today, and what they cannot

The pitch from both companies is "no rewrite". The reality, based on their own documentation, is "no rewrite for many common cases".

Likely to move with little effort

Standard transformer models loaded through Hugging Face libraries, which both companies now support upstream;
training jobs built on TorchTitan or common FSDP-style sharding;
LLM serving through vLLM, which both clouds now support;
code that already uses PyTorch's device-neutral APIs instead of hard-coding cuda.

Still hard, or not supported yet

Custom CUDA kernels. Hand-written Nvidia kernels do not run on other chips. Teams must rely on vendor kernel libraries or rewrite them in NKI (Amazon) or Pallas (Google).
Dynamic shapes. Amazon lists "dynamic and symbolic shapes" among features TorchNeuron chose not to support in its first release. Google's engineers wrote that changing sequence lengths and batch sizes trigger recompilation on TPUs, and listed better dynamic-shape support as a roadmap item.
Less common operations. Amazon's list of unsupported features includes sparse tensors, NestedTensor, forward-mode autograd and some upsampling modes. In eager mode, an operation without a Neuron version runs on the CPU with a warning, which can slow a model sharply.
CUDA-specific tooling. Features such as CUDA graphs and NCCL-specific collectives are, by definition, Nvidia-only.

"Runs" is also not the same as "runs fast". Both companies stress that compiled mode and tuned kernels are needed for production-level speed. Customer claims vary by workload. Deepgram's chief technology officer, Adam Sypniewski, said in Amazon's post that Trainium processed about 60% more input audio seconds per dollar than the company's existing GPU baselines "in the benchmarks we evaluated". That is a single customer's measurement and should not be read as a general result.

Why the software layer matters more than chip specs

Raw chip performance numbers are easy to publish and compare. Switching costs are harder to see. A company's real investment sits in its model code, training recipes, debugging habits, monitoring tools and the engineers who know them. Amazon's PyTorch post makes this point directly: the cost of running AI includes the engineering time needed to get workloads working and keep them working as libraries change.

Nvidia's position has rested on making that engineering time smallest on its own hardware. As TechCrunch notes, many PyTorch features were built to work with Nvidia GPUs because those chips were already the standard for AI research. A rival chip with better specifications can still lose if moving to it costs months of porting.

Native backends attack that cost directly. They do not have to beat Nvidia on every benchmark. They only have to make the switch cheap enough that price, availability or capacity can tip the decision. For Amazon, the stakes are large: it says its custom silicon business has passed an annual revenue run rate of $25 billion. Google, for its part, offers TPUs to outside customers through Google Cloud. Meta says TorchTPU lets it deploy its own PyTorch workloads on those chips.

There is also a quieter shift inside PyTorch itself. Each upstream change that removes a hidden CUDA assumption helps every non-Nvidia chip, not just Amazon's or Google's. Meta's director of PyTorch engineering, Jana van Greunen, is quoted in both companies' announcements. In Amazon's post she said developers "should have a consistent experience across hardware platforms".

What it means

For AI developers, the practical effect is more choice. Teams can more easily try the same PyTorch model on Nvidia GPUs, Trainium or TPUs and compare cost and speed on their own workloads. The best first test is a standard model on a supported library, not a heavily customised research codebase.

For Nvidia, this week does not end its software lead. TechCrunch noted that developers will not automatically leave Nvidia hardware. Nvidia's tools are mature, and custom kernels written for CUDA still have to be ported by hand. Neither Amazon's nor Google's tool is a general replacement yet. Amazon calls its release a beta, and Google labels parts of its stack as previews.

What has changed is the direction. Two hyperscalers now treat PyTorch, not their own proprietary frameworks, as the front door to their chips. Both also say they are contributing fixes back to PyTorch itself. If those efforts continue, the key question for AI teams may shift from "will my code run on this chip?" to "which chip runs my code most cheaply?" That shift would put pressure on prices and margins across the AI hardware market. Several things are still unknown: how quickly the remaining gaps close, how independent benchmarks compare, and how Nvidia responds.

The question of cheaper compute is becoming more pressing as AI developers raise ever larger sums. In recent days Bullquill has covered reports that OpenAI is nearing a $30 billion raise at a $1.25 trillion valuation and that DeepSeek has reportedly raised $15 billion with backing from Tencent and CATL, while established vendors such as IBM, whose software revenue rose 8% in the third quarter, are competing for the same enterprise AI budgets.

Frequently asked questions

What is a PyTorch backend?

A backend is the layer that turns PyTorch operations into work a specific chip can run. TorchTPU and TorchNeuron register Google TPUs and AWS Trainium through PrivateUse1, PyTorch's official extension point for new hardware, so developers can target them with a device setting.

Can I run any PyTorch model on TPUs or Trainium without changes?

Many standard models, especially those loaded through Hugging Face libraries or trained with TorchTitan, can move with minimal changes. Custom CUDA kernels, dynamic input shapes and some less common operations are still unsupported or limited, according to the companies' own documentation.

Is TorchNeuron open source?

Yes. Amazon released TorchNeuron under the Apache 2.0 licence on GitHub (aws-neuron/torch-neuronx), and it entered public beta on 20 October 2026 for Trainium2 and Trainium3.

Does this end Nvidia's CUDA advantage?

Not yet. Both tools are new, parts are labelled beta or preview, and Nvidia's software tools remain mature. The releases lower switching costs, which makes it easier for teams to compare chips on price and performance.