Amazon Web Services Announces TorchNeuron for PyTorch on Trainium Chips

What is new

Amazon Web Services (AWS) has announced TorchNeuron, a new PyTorch plugin designed to allow AI developers to run models on Amazon’s Trainium chips [1, 2]. The tool provides a native backend for PyTorch, enabling support for eager mode, torch.compile, and distributed APIs such as FSDP and TP [2].

Key details regarding the release include:

• TorchNeuron will be available as a public beta starting October 20, 2026 [2].
• Developers can install the tool using the command pip install torch-neuronx [2].
• The tool is designed to work with the PyTorch framework, which Amazon states is used for approximately 85% of AI applications [1].
• Amazon plans to contribute code to the main PyTorch project to make the framework more hardware-agnostic and less tied to specific chips [1].

According to AWS, the tool allows apps written in PyTorch to run on Trainium by changing only one line of configuration, without requiring a code rewrite [1]. For Hugging Face users, the tool provides native support for Transformers, Diffusers, Accelerate, PEFT, and TRL libraries [2]. In terms of text generation, TorchNeuron supports 89% of the top 100 most-used model architectures out of the box, an increase from the 46% supported by the previous Optimum Neuron v0.4.5 [2].

Deepgram, an AI voice application provider, has already used this technology to port its Nova-3 model to AWS Trainium2 [3]. Deepgram reported that their setup processed between 11% and 70% more audio per dollar on offline, throughput-oriented batch workloads compared to their previous compute options [3].

How to use it

For developers using Hugging Face libraries, the process for moving models to Neuron devices involves the following:

• For inference, users can pass the device_map=neuron argument in their code [2].
• For training, users can use standard libraries like Transformers, Accelerate, PEFT, and TRL [2].
• To enable torch.compile, users can pass the --torch_compile true flag; on a Trainium instance, the system will automatically pick the Neuron backend [2].
• When running training via torchrun, users should set the TP_SIZE and FSDP_SIZE to match the number of processes, such as 32 processes for a 4 x 8 grid [2].
• Users must currently disable built-in asynchronous weight loading in Transformers by setting the environment variable HF_DEACTIVATE_ASYNC_LOAD=1 [2].

Limits, privacy and academic integrity

There are specific technical constraints and requirements when using TorchNeuron:

• The tool requires the use of static shapes. The script disables automatic dynamic shapes in torch.compile and uses wrapped packing for sequences up to 2048 tokens [2].
• With TRL version 1.16 or higher, batches on Neuron are padded to a multiple of 128 tokens to limit the number of input shapes [2].
• Asynchronous weight loading in Transformers is not supported on Neuron at this time [2].
• While AWS claims the tool improves economics, they noted that AWS may not offer the cheapest per-token price [1].