Amazon's TorchNeuron, now in public beta, runs PyTorch code on Trainium chips as a standard device

Amazon has released TorchNeuron, software that lets PyTorch run directly on Trainium, the AI chips it rents out through AWS. Announced on October 19, 2026, it became a public beta on October 20, 2026, under the Apache 2.0 license. PyTorch is the most widely used framework for building and training AI models, and earlier versions of Amazon's Neuron software reached Trainium through a separate layer called PyTorch/XLA. With TorchNeuron, the chip appears to PyTorch as an ordinary device called neuron.

What changes in the code

A backend is the layer that tells PyTorch how to run its calculations on a specific piece of hardware. TorchNeuron registers through PrivateUse1, the official extension point PyTorch provides for that purpose. Moving a model to the chip takes .to('neuron'), and compiled models name the backend with torch.compile(model, backend="neuron").

The supported modes are eager execution, where each operation runs immediately, and torch.compile, which turns the model into an optimized graph. Developers can also write custom low-level kernels through the Neuron Kernel Interface (NKI). An adaptive eager mode, meant to bring eager code closer to compiled speed, is labeled experimental.

The surrounding libraries were updated as well. Trainium is a supported device in Hugging Face Transformers, Diffusers, and TRL, so a model loads onto the chip with device_map="neuron". Distributed training works through standard PyTorch tools such as FSDP2 and TorchTitan.

What the first release does not cover

Installation requires a Trn2 or Trn3 instance, Python 3.12, 3.13, or 3.14, and PyTorch 2.14. Models can use BF16, FP32, and FP16. FP8 and the TorchAO quantization library are still limited to a private beta.

AWS lists features it deliberately left out of the first release, including sparse tensors, NestedTensor, forward-mode autograd, and dynamic and symbolic shapes. The team says customers had not asked for them.

The test results show how far coverage reaches. In an August 2026 run of PyTorch's own test suites, about 154,000 test outcomes were recorded. About 39,000 tests were deliberately skipped for Neuron, roughly 22,000 of them because the operations have no native Trainium implementation yet.

Claims nobody has measured independently

Amazon says existing PyTorch code runs on Trainium with a single configuration change and no rewrite. The savings figures come from Amazon's customers. Deepgram, a speech recognition company, says that in benchmarks it chose, Trainium processes about 60% more seconds of input audio per dollar than its existing GPU setups.

No independent comparison of TorchNeuron's speed or cost against other chips has been published, and the announcement does not give instance prices. This suggests the practical first step for a team is narrower than a migration: checking whether its model relies on any of the features the beta does not support.