TorchNeuron brings native PyTorch support to AWS Trainium chips

A standard PyTorch device for Trainium

According to the PyTorch announcement, TorchNeuron is meant to ensure that the model code, libraries and tools developers already use run on AWS Trainium. Trainium is described as a portfolio of AI chips designed for cloud scale AI training and inference. The backend supports Trainium2 and Trainium3.

The integration builds on accelerator-agnostic extension points that the PyTorch community created, including the PrivateUse1 device backend and the torch.accelerator APIs. TorchNeuron registers Trainium through PrivateUse1, so standard APIs such as torch.nn, torch.tensor, torch.autocast, torch.distributed and torch.profiler operate on the neuron device.

Switching hardware is meant to be a small change. In the example given, a model is moved to the chip with .to('neuron'), and the rest of a training loop stays the same. Code written against torch.accelerator, PyTorch's recommended device-agnostic API, runs unmodified. For compiled models, developers name the Neuron backend once with torch.compiler.set_default_backend("neuron") or per call with torch.compile(model, backend="neuron").

Jana van Greunen, Director of PyTorch Engineering at Meta, said in the announcement that "developers should have a consistent experience across hardware platforms" and that teams can use Trainium's capabilities "without reinventing their workflows or leaving the PyTorch ecosystem."

Eager mode, torch.compile and distributed training

TorchNeuron supports both of PyTorch's execution modes. Eager mode runs operations immediately, so Python control flow, prints and breakpoints work without a compile step. The post positions it for research and debugging. Because eager mode dispatches operations individually, each one compiles its own Neuron Executable File Format (NEFF) file, launches on the device and moves intermediates through HBM.

To reduce that cost, the team is introducing an experimental adaptive eager mode. According to the post, it lowers launch overhead and offers more chances to fuse operations, approaching the performance of a compiled graph without code changes while keeping the interactive experience.

For production, torch.compile captures the model into a graph that the Neuron compiler optimizes. The pipeline uses TorchDynamo to capture an FX graph, AOTAutograd to trace the backward graph, and Torch-MLIR to lower the result into StableHLO before the Neuron compiler produces a NEFF file. FlexAttention is supported natively. The Neuron inference stacks use the same support: vLLM Neuron serves LLMs, and vLLM Omni Neuron serves diffusion models for image and video generation.

On data types, models can use BF16, FP32 and FP16, along with integer and boolean types. TorchAO and FP8 are in private Beta and are expected to become generally available in the next release.

Distributed workloads can use standard PyTorch APIs, including DDP, FSDP2, SimpleFSDP, tensor parallelism and pipeline parallelism, plus compositions such as FSDP2 with tensor parallelism and ZeRO. DTensor, DeviceMesh, collective APIs such as all_gather, reduce_scatter and all_reduce, and the P2P APIs send and recv are supported. Collectives run over Trainium's dedicated collective cores within a node and over the Elastic Fabric Adapter (EFA) between nodes.

Streams, custom kernels and profiling

Torch streams and events are supported through torch_neuronx.Stream and torch_neuronx.Event, which map PyTorch's stream semantics onto Trainium's hardware queues. The post notes some limits: non-default streams multiplex onto a single hardware queue, so they keep ordering, priority and event semantics but do not each get a dedicated engine. Errors surface at the next synchronization point, where TorchNeuron re-raises the earliest exception with its original context.

The Neuron Kernel Interface (NKI) lets developers write, optimize and run custom operators in both eager and torch.compile. Researchers can use it to prototype ideas needing low level hardware control, including new model architectures. Developers can call performance critical NKI kernels directly to replace sequences of standard ATen ops.

TorchNeuron also ships NKI-lib, an open source set of Trainium optimized operators for dense and MoE models, including FlashAttention and MoE blockwise kernels. For profiling, Neuron Explorer works with the standard torch.profiler to trace execution from framework to hardware and link events back to source lines.

What it means for developers

For teams already building on PyTorch, TorchNeuron lowers the effort of trying Trainium: existing code, distributed strategies and profiling tools carry over, with custom kernels available where more control is needed. The code is to be published at github.com/aws-neuron/torch-neuronx as a public Beta, and the announcement came ahead of the open source launch at PyTorch Conference North America on Oct. 20. Some features, including adaptive eager mode and TorchAO with FP8, are still experimental or in private Beta.