Google Open-Sources TorchTPU as Amazon's TorchNeuron Enters Beta

Google open-sourced TorchTPU on Tuesday 20 October, the opening day of the PyTorch Conference in San Jose, making its Tensor Processing Units a native target for PyTorch code. Amazon's TorchNeuron, announced the day before, entered public beta the same day for Trainium2 and Trainium3. In Google's version a model moves to a TPU by changing the device string; in Amazon's, the edit in the company's example is model.to('neuron'). Both are on GitHub, and both address the same cost: the engineering time a team spends making a working PyTorch codebase run somewhere other than Nvidia hardware.

Related reading: Amazon Looks to Offload $8 Billion of Nvidia Chips Through a Financing Vehicle · DeepSeek Open-Sources a Software Stack for Huawei's Ascend Chips, Aiming at CUDA.

Neither is a new chip. Each is a software backend that sits between PyTorch and an accelerator, intercepts the framework's operations, and lowers them into something the chip can execute. TorchTPU targets TPUs and TorchNeuron targets Trainium. Both attach through PyTorch's PrivateUse1 extension point, the mechanism that lets a hardware vendor register a device without forking the framework. The extension point is why two rival clouds could ship native backends in the same week. Neither of them owns PyTorch.

What the two releases actually contain

TorchTPU is the older project. Google says work began in mid-2025 with Meta's PyTorch team, and the PyTorch blog calls it "an open source collaboration between Google and Meta that introduces a native and principled PyTorch backend for Google TPUs." It replaces the earlier route through PyTorch/XLA, which Google introduced in 2020 on lazy tensors. That version required developers to mark compilable regions in their source, had long compile times, and did not support PyTorch's eager execution mode. TorchTPU runs eagerly by default, supports torch.compile, registers about 450 Aten operators, and implements TPU collectives at the process-group level so distributed code written for DTensor, FSDP and similar APIs works without a separate communication library.

Google's announcement lists four repositories: TorchTPU itself, plus TPU backends for vLLM and SGLang and a TorchTitan integration. The serving engines matter more than the headline. vLLM and SGLang are two of the main engines for serving open-weight models, and a backend that only worked inside a training script would leave inference on GPUs. Google also names integrations with llm-d 1.0, Ray, TorchRec, Hugging Face Transformers v5, where device_map="tpu" loads a model onto a TPU, and LiteRT, its on-device runtime, which joined the PyTorch ecosystem. A TPU agent suite in preview is meant to answer migration questions and fix errors during a port.

The coverage claim is concrete; the performance claim is not. Google says it validated 4,917 of the most downloaded Hugging Face models across more than 500 architecture families, and the PyTorch post says TorchTPU is already used in production for inference and training by "some of the largest Google Cloud customers," without naming them. Running a model out of the box and running it well are different results, and Google's own post says peak performance requires shapes, sharding and operators aligned with the TPU's architecture.

“With TorchTPU, TPUs join GPUs and CPUs as a native option for PyTorch.” — Mark Lohmeyer and Bill Jia, Google Cloud, 20 October 2026

Amazon's announcement is narrower, and its documentation is more explicit about the gaps. TorchNeuron registers Trainium as a neuron device, supports eager and torch.compile, exposes custom-kernel tooling through the Neuron Kernel Interface, and sends profiling data through the standard torch.profiler into a viewer called Neuron Explorer. It supports DDP, FSDP2, tensor and pipeline parallelism, DTensor and DeviceMesh. Hugging Face added Trainium support upstream in Transformers, Diffusers, TRL and its kernels library, so device_map="neuron" works there as well. The package is pip install torch-neuronx, and the code is Apache 2.0 at aws-neuron/torch-neuronx.

The gaps are documented by Amazon, not discovered by critics. The first release does not support dynamic or symbolic shapes, sparse tensors, forward-mode autograd or several smaller features. In an August 2026 run of PyTorch's upstream test suites, Amazon recorded about 154,000 test outcomes. Roughly 69,000 were skipped for a valid reason and about 39,000 were deliberate Neuron skips; the largest group, around 22,000, were operators without a native Neuron implementation. TNW summarised the same material as roughly 22,000 PyTorch tests skipped. Both are accurate: one counts the whole suite, the other only the skips that point to missing operator coverage.

One line of code, and the months around it

The device string is the demo, not the migration. Google's post says a one-line change gets a model running, then adds the caveat about peak performance. Amazon's demonstration ran the same training script on Trainium and a conventional chip and said the model learned "the same things at the same rate"; no throughput, cost or convergence numbers were published.

What changed is the direction of the work. Before these backends, a port to TPU or Trainium meant adapting to a vendor-specific stack, rewriting kernels, revalidating numerics and maintaining a fork against upstream PyTorch. Now the first attempt is a configuration change and the remaining work is tuning. That is a smaller bill at a time when buyers are watching costs. Amazon talks about engineering time. Google talks about choice across TPUs, GPUs and CPUs under one operational model.

The TechCrunch report on the pair put the strategic reading plainly: "Amazon and Google are clearly working to end Nvidia's advantage." The same article added that none of this means developers will ditch Nvidia's hardware en masse. A vendor can remove a barrier without winning the customer.

Meta's involvement points at optionality. Jana van Greunen, director of PyTorch engineering at Meta, said TorchTPU lets Meta "deploy our PyTorch workloads to Google TPUs, provide direct feedback to the engineering teams, and build the ecosystem flexibility everyone needs." That describes a second path to production, not a departure from Nvidia. Meta co-developed the backend.

The commercial stakes are large enough to explain the coordination. Amazon says its custom chip business passed a $25 billion annual revenue run rate, and chief executive Andy Jassy said in April 2026 that it could be worth $50 billion (TNW). Google announced its eighth-generation TPUs in April 2026, with the 8t for training and 8i for inference, and has an expanded Anthropic deal for multiple gigawatts of TPU capacity. According to Yahoo Finance, it is also working to supply OpenAI. Every workload that runs on those chips without a rewrite is capacity the cloud can fill with silicon it controls.

The benchmarks are vendor benchmarks

Google's headline number is real but narrow. It says TorchTitan, its training reference stack, reached up to 52% model FLOPs utilisation on Qwen 3 8B while scaling to 8,192 Ironwood TPUs, and that the new DeferAndFuse eager mode can be 50% or more faster than strict eager execution. Both are Google measurements, on Google hardware, with the company's software configuration. They are evidence that the software stack is less of a bottleneck than it was. They do not show that a TPU is faster or cheaper than a GPU for a given job.

Amazon's most useful number comes from a customer rather than a suite. Deepgram, which sells speech models, is quoted saying Trainium processes "about 60% more input audio seconds per dollar than our existing GPU baselines" after a multi-month port. The baseline is Deepgram's own fleet, not a named Nvidia part, and the measure is audio seconds per dollar on a speech workload. That is meaningful for anyone serving speech and not a general result.

Why a backend does not undo CUDA

A backend removes one layer of switching cost. It does not remove the others. CUDA's advantage is not a single API that PyTorch can abstract away. It is the accumulated set of libraries, tuned kernels, profilers, debuggers and answers to other people's problems that have built up around Nvidia's hardware for more than a decade. A team with a custom attention kernel, a fused optimiser and a year of performance work has assets that a device string does not carry across.

The launch material concedes this in places. Google's TorchTPU post says its default collective implementation compiles each collective into a separate module, which can cost 30% to 40% end to end, and that bounded dynamic shapes are still being built; torch.compile currently forces dynamic=False. Amazon's first TorchNeuron release skips dynamic shapes and the other features listed above. Open source does not mean finished. TorchNeuron is a public beta, and Google's documentation describes maturity that varies by feature.

PyTorch itself is neutral ground now. Meta wrote the framework in 2016 and transferred it to the Linux Foundation in 2022, so a Google or Amazon backend can arrive as ecosystem work rather than a proprietary detour. The foundation's conference is where both releases landed. That arrangement gives the vendors a reason to improve the shared layer: a fix that lands in PyTorch helps their hardware without asking developers to adopt a vendor stack.

The more consequential work may be in the pull requests. Amazon says it is submitting changes to mainline PyTorch to remove hard-coded CUDA assumptions, and the TorchNeuron post lists them. Google contributed LiteRT to the PyTorch ecosystem and worked through Hugging Face, vLLM and SGLang rather than around them. Those changes are what would make the next port cheaper.

What would count as evidence

An independent benchmark would settle more than either launch post. The test is easy to specify: the same model, batch sizes and precision, measured on a TPU, a Trainium instance and an Nvidia GPU, with cost per token included. A named customer describing a completed migration and the bill before and after would carry similar weight. So would merged upstream code: operator coverage and device-agnostic changes accepted into PyTorch rather than sitting in a vendor repository.

Caveats

The two releases are software announcements, not performance results. Neither company published independent, like-for-like performance numbers for either stack on a common model against Nvidia hardware. TorchNeuron is a public beta with documented feature gaps, and TorchTPU's documentation describes varying maturity, a stated penalty on its default collective path and unfinished work on dynamic shapes. The "85%" figure that appears in coverage is not one number: Amazon says 85% of AI researchers and developers write in PyTorch, while TechCrunch says PyTorch is used by 85% of AI apps "by some accounts". Both are attributed, and they measure different populations. Migration will also be gradual: a backend lowers the cost of trying a second supplier, and procurement still runs on measured cost per token and on delivery dates.

Both backends can be downloaded now. The benchmark that would settle what the porting is worth has not been published yet.

Cover image: "Google Data Center, Council Bluffs Iowa" by Chad Davis, CC BY 2.0, via Wikimedia Commons.