Overview
TorchNeuron is a native PyTorch backend for AWS Trainium chips, released by AWS on October 19, 2026, and opened as a public Beta under the Apache-2.0 license at PyTorch Conference North America on October 20, 2026. Its code repository is located at github.com/aws-neuron/torch-neuronx. This tool adds a neuron device to upstream PyTorch through PyTorch's PrivateUse1 mechanism, so developers only need to move models to neuron for training or inference while continuing to use existing PyTorch APIs. TorchNeuron uses eager mode by default, executing operators one by one on the device without first tracing a computation graph; when performance is needed, the neuron backend of torch.compile can be used to compile the model or part of the model. It is a different technical path from the older XLA-based torch-neuronx, and this repository is positioned as a PyTorch native backend.
Key Features
- Native PyTorch device integration: Adds a neuron device to upstream PyTorch through PyTorch's PrivateUse1 mechanism, so models can be moved to neuron for training or inference while continuing to use existing PyTorch APIs.
- Default eager execution mode: Uses eager mode by default, executing operators one by one on the device without first tracing a computation graph, which facilitates debugging and rapid iteration.
- torch.compile acceleration: When performance is needed, the neuron backend of torch.compile can be used to compile the model or part of the model, improving execution efficiency while preserving the PyTorch programming experience.
- Distributed training support: Supports torch.distributed, DTensor, and functional collectives, covering distributed training scenarios such as data parallelism, tensor parallelism, and fully sharded parallelism.
- NKI custom kernel invocation: NKI custom kernels can be called from PyTorch code and are available in both eager mode and within torch.compile, making it easy to perform low-level optimization for specific operators.
- Preinstalled runtime environment: Neuron DLAMI and Deep Learning Container come with TorchNeuron preinstalled; other Linux instances require the Neuron driver and runtime to be installed first. The runtime environment requires Trn2 or Trn3 instances and Python 3.12/3.13/3.14.
Use Cases
- Model training on AWS Trainium chips using native PyTorch APIs
- Model inference on AWS Trainium chips using native PyTorch APIs
- Compiling the model or part of the model with the neuron backend of torch.compile to improve performance
- Data parallel, tensor parallel, and fully sharded parallel training based on torch.distributed, DTensor, and functional collectives
- Calling NKI custom kernels from PyTorch code for operator-level optimization
Pros
- Open source under the Apache-2.0 license, with a publicly accessible code repository
- Native PyTorch device integration through the PrivateUse1 mechanism, continuing to use existing PyTorch APIs with low migration cost
- Default eager mode requires no prior computation graph tracing, facilitating debugging and rapid validation
- Supports the neuron backend of torch.compile, allowing the model or part of the model to be compiled as needed
- Distributed training covers data parallelism, tensor parallelism, and fully sharded parallelism, and supports torch.distributed, DTensor, and functional collectives
- NKI custom kernels can be called from PyTorch code and are available in both eager mode and within torch.compile
Pricing
TorchNeuron is open source under Apache-2.0 with no separate license fee; the cost of use is the Trn2 or Trn3 instance plus related AWS services. Neuron DLAMI and Deep Learning Containers ship with it preinstalled, while other Linux instances need the Neuron driver and runtime installed first. Refer to AWS pricing for instance rates.
Summary
TorchNeuron is a native PyTorch backend launched by AWS for Trainium chips, opened as a public Beta under the Apache-2.0 license. It adds a neuron device through the PrivateUse1 mechanism, uses eager execution by default and supports torch.compile compilation, covers data parallelism, tensor parallelism, and fully sharded parallelism for distributed training, and can also call NKI custom kernels. The runtime environment requires Trn2 or Trn3 instances and specified Python versions, and it is a different path from the older XLA-based torch-neuronx.
Version History
- AWS releases TorchNeuron, a native PyTorch backend for Trainium (2026-10-19): AWS released TorchNeuron on 2026-10-19, a native PyTorch backend for AWS Trainium under Apache-2.0, entering public beta at PyTorch Conference North America on 2026-10-20. It adds a neuron device through PyTorch PrivateUse1, runs eagerly by default, supports torch.compile, covers distributed training through torch.distributed, DTensor and functional collectives, and can call NKI custom kernels. It requires Trn2 or Trn3 instances and Python 3.12 to 3.14.