Amazon wants PyTorch code to run on its own Trainium chips

Amazon Web Services (AWS) has announced TorchNeuron, an open source extension to the widely used AI framework PyTorch. It is designed to let developers run existing model code, libraries and tools on Amazon's own AI chips, Trainium. According to the plan, it launches as a public beta on October 20, 2026, during the PyTorch Conference North America.

A small change in the code

According to AWS, in the simplest case the switch means changing the device to "neuron," for example with .to('neuron'), instead of CUDA, which is Nvidia's system. Compiled models require specifying the Neuron backend in torch.compile. Installation is done with pip install torch-neuronx, and the code is released under the Apache 2.0 license on GitHub.

The backend supports Trainium2 and Trainium3. It can also be used with Hugging Face Transformers v5, Diffusers, TRL and TorchTitan, and there are vLLM plugins for inference, meaning running finished models. AWS says it continuously runs 100,000 Hugging Face models on Trainium, which AWS says covers 99.9% of real-world usage measured by downloads.

What do users say?

Deepgram says that in the benchmarks it evaluated, Trainium processes roughly 60% more seconds of input audio per dollar than its current GPU solutions. That is the company's own figure and has not been independently verified. Porting the model took several months with help from AWS.

Cartesia co-founder Arjun Desai says they ran their standard pipeline with minimal changes to the PyTorch code, but that they write NKI kernels themselves for the layers that drive latency. Hugging Face calls TorchNeuron a practical realization of its "silicon neutrality."

Gaps in the first release

The first release does not support dynamic shapes, sparse tensors, NestedTensor and forward-mode autograd, among other things. TorchAO and FP8 are only in private beta and are expected to be generally available in the next release. In PyTorch's own tests from August 2026, about 69,000 of roughly 154,000 test outcomes were skipped. Of those, roughly 22,000 lack a native Neuron lowering of operators.

According to AWS's documentation, operators without support fall back to the CPU in eager mode with warnings. AWS says it is continuing to work on closing the gaps.

The Next Web points out that the release has significant gaps and that the tool targets Nvidia's position of strength. The outlet names Cartesia, Deepgram, Lieber Institute, Hugging Face and Uber as early users. It also writes that Amazon's chip division has roughly $25 billion in annual revenue. That figure is not confirmed in AWS's own text.

For ordinary AI developers, the takeaway is that it may eventually become easier to choose chips other than Nvidia's without rewriting everything. How easy it becomes in practice depends on how many of the gaps get closed.