Amazon TorchNeuron Lets PyTorch Code Run on Trainium Chips

Amazon has released TorchNeuron, an open-source software layer that lets PyTorch code run directly on AWS’s Trainium2 and Trainium3 AI chips. The public beta opened on October 20, 2026, at PyTorch Conference North America. Amazon says moving over from another chip takes a single device setting.

What Amazon and PyTorch announced

PyTorch is the framework used to build and train AI models. Trainium is an AI chip designed by AWS and rented through Amazon’s cloud. The two have not fit together automatically: according to Techzine, Amazon says porting PyTorch code to other chips used to mean months of rewriting and revalidation.

TorchNeuron is meant to remove that work. Amazon announced it on October 19, 2026 and opened the public beta a day later. The code is licensed under Apache 2.0 and hosted on GitHub as aws-neuron/torch-neuronx. Developers can install it with pip install torch-neuronx. It supports Trainium2 and Trainium3.

Amazon says about 85 percent of AI researchers and developers write their code in PyTorch, which makes native support the condition for Trainium to be considered at all. The Next Web places the launch within Amazon’s push to compete with Nvidia, whose chips and CUDA software dominate AI computing. The chip race also showed up in Tenstorrent’s funding for open AI chips.

The PyTorch team at Meta built the extension mechanism that TorchNeuron relies on, and Meta contributed a statement to the PyTorch blog. Jana van Greunen, Director of PyTorch Engineering at Meta, wrote that “developers should have a consistent experience across hardware platforms.”

How it works and what the beta lacks

TorchNeuron registers Trainium as a standard PyTorch device through an official extension point called PrivateUse1. Developers move a model or tensor to the neuron device, for example with .to('neuron'), and standard PyTorch features such as distributed training and the profiler keep working.

There are two execution modes. Eager mode runs operations one at a time and suits research and debugging. torch.compile with the neuron backend turns the model into a graph optimized by the Neuron compiler, and is meant for stable production models. Where more control is needed, developers can write custom kernels with the Neuron Kernel Interface (NKI). AWS also released NKI-lib, a kernel library that includes FlashAttention and mixture-of-experts kernels.

TorchNeuron supports PyTorch’s standard distributed APIs, including DDP, FSDP2, tensor parallelism, and pipeline parallelism, and it integrates with TorchTitan, PyTorch’s library for large-scale training. The built-in PyTorch profiler can trace bottlenecks down to the source line through the Neuron Explorer tool. Amazon also says it is contributing changes to PyTorch and vLLM so that framework code depends less on specific chips.

The first release is a beta with gaps that AWS acknowledges on the PyTorch blog. In an August 2026 test run, about 22,000 tests from the PyTorch suite were skipped because the operators involved had no native Neuron implementation. Features deliberately left out of the first release include sparse tensors, NestedTensor, forward-mode autograd, and dynamic input shapes. FP8 and TorchAO remain in private beta.

The Neuron documentation spells out the practical effect. In eager mode, operators without support run on the CPU with a warning, so the model still runs but may slow down. Under torch.compile, unsupported constructs break the graph and reduce optimization opportunities. The float64 type is also executed as float32.

AWS says it routinely tests TorchNeuron against 100,000 Hugging Face models that cover 99.9 percent of real-world usage by download count. By its own account, results as of September 2026 still leave gaps that it plans to close in later releases.

Early users and the cost claims

The PyTorch blog says TorchNeuron is already used in autonomous vehicles, robotics, video generation, speech, telecommunications, healthcare, and geospatial modeling. Three examples are worth comparing.

In an internal test, Amazon ran the same training script on a conventional chip and on Trainium, changing only the device setting, and said the model learned the same things at the same rate on both. That result has not been verified independently.

Cartesia, a voice AI developer, ran its standard inference pipeline on Trainium with minimal changes to its PyTorch code. Its team then focused on writing NKI kernels for the layers called most often. Deepgram, a speech recognition company, reported something more concrete: in the benchmarks it evaluated, Trainium processed about 60 percent more audio seconds per dollar than its GPU baselines. That figure came from a multi-month porting project with the AWS Generative AI Innovation Center, not from a quick switch.

The Lieber Institute for Brain Development, a research institute that studies conditions including schizophrenia, uses TorchNeuron to explore AI models that could speed up or replace computationally heavy molecular dynamics simulations, according to Amazon. Hugging Face has also added Trainium support to its Transformers and Diffusers libraries, so open models can be loaded onto the chip without device-specific code.

All performance figures above come from AWS or its customers. Techzine notes that the launch came without independent benchmarks or pricing details, so cost comparisons need to be tested in-house.

What this means for data and AI teams

For banks, hospitals, or agencies that train and run their own models in the cloud, TorchNeuron adds a hardware option without forcing a change of framework. It is worth weighing, but through a measured pilot.

A few concrete points for a pilot:

1. Test your own model. The “single setting” claim describes a starting point. Models with dynamic input shapes or rare operators may need adjustments.
2. Read the CPU fallback warnings. A model can run without errors while part of its work quietly lands on the CPU.
3. Compare cost per job. An hourly chip price is not the cost of finishing a training run. Measure total time and cost for the same workload on GPUs and on Trainium.
4. Check regions and data location. The launch sources do not list the AWS regions that offer Trainium. Confirm availability and data residency rules against your organization’s policy, including personal data protection requirements.
5. Watch for new lock-in. Model code stays standard PyTorch, but NKI kernels run only on Trainium. The more custom kernels you write, the costlier it becomes to move again.

Organizations without an in-house ML team do not need to rush. During a beta, chip choice usually follows the workloads that already exist. Merakit helps organizations map those workloads through its machine learning, cloud computing, and artificial intelligence services.

Frequently asked questions

What is TorchNeuron?

TorchNeuron is an open-source PyTorch backend from Amazon that registers Trainium chips as a neuron device. PyTorch models can then run on Trainium2 and Trainium3 using familiar PyTorch functions.

Does PyTorch code run unchanged?

For simple cases, Amazon says changing the device setting is enough. In practice the beta does not yet support dynamic input shapes and some operators, and production performance usually calls for torch.compile. Cartesia itself reported only minimal changes to its code, but not zero.

Is TorchNeuron free, and which chips does it support?

The software is free under Apache 2.0 and supports Trainium2 and Trainium3. Renting Trainium chips on AWS is still billed, and the launch sources do not discuss pricing.

Conclusion

TorchNeuron lowers the technical barrier to trying Trainium, but its real value will show only once many teams test it on their own models. Start with one model, measure the result, and decide on data. The Merakit team is ready to talk if your organization wants to plan such a pilot.

Sources

• Introducing TorchNeuron: Native PyTorch backend for AWS Trainium, PyTorch Blog (updated October 20, 2026)
• Amazon’s TorchNeuron lets PyTorch code run on AWS Trainium chips, The Next Web (October 19, 2026)
• TorchNeuron brings PyTorch code to AWS Trainium without modifications, Techzine (October 20, 2026)
• TorchNeuron and PyTorch on AWS Trainium chips, About Amazon (October 2026)
• AWS Neuron documentation: PyTorch native overview