Introducing TorchNeuron: Native PyTorch backend for AWS Trainium

Frontier models span language, video and audio, on architectures from dense and mixture-of-experts transformers to diffusion, state-space, world and structure-prediction models. PyTorch is where AI research happens, where the models are defined, and increasingly where production happens too. Researchers want to experiment and innovate, and customers want the performance AWS Trainium offers to enable them to achieve their performance and scale goals. The training and serving libraries are written against PyTorch’s own APIs, and new libraries and new versions ship continuously. Over the last several years, PyTorch made those APIs accelerator-agnostic, with official extension points for registering custom hardware. Today, to address the growing demand from the developers and research community on Trainium, we are excited to share an even simpler approach to program the chip as a standard PyTorch device, ahead of its open source launch at PyTorch Conference North America on Oct. 20. AWS Trainium is a portfolio of AI chips designed from the ground up for next generation cloud scale AI training and inference. We are announcing TorchNeuron , an open source native PyTorch backend for Trainium, so the model code, the libraries, and the tools you already use run on Trainium. This integration builds directly on the accelerator-agnostic extension points, including the PrivateUse1 device backend and torch.accelerator APIs, created by the PyTorch community. TorchNeuron supports eager execution for research and debugging including new adaptive eager mode with improved performance, torch.compile for production performance, and the Neuron Kernel Interface (NKI) for custom kernels, on Trainium2 and Trainium3. Also, with Neuron Explorer , performance engineers use the same torch.profiler to trace execution from framework to hardware, linking every event back to its source line to pinpoint bottlenecks and optimize efficiently. TorchNeuron will be released as open source under Apache 2.0 at github.com/aws-neuron/torch-neuronx , going live as public Beta on Oct. 20. “PyTorch’s vision is straightforward: developers should have a consistent experience across hardware platforms. Native support for AWS Trainium expands that choice, enabling researchers and engineers to experiment and iterate using the eager execution and distributed APIs they already know. Teams can take advantage of Trainium’s capabilities without reinventing their workflows or leaving the PyTorch ecosystem.” – Jana van Greunen, Director of PyTorch Engineering at Meta Getting started: Easy device switching Getting started is straightforward. Run your code on Trainium using the pointer: .to('neuron') . # Run a training loop on Trainium model = MyModel().to('neuron') # the only change optimizer = torch.optim.AdamW(model.parameters()) for batch in dataloader: optimizer.zero_grad() loss = criterion(model(batch), targets) loss.backward() optimizer.step() Modern code written against torch.accelerator, PyTorch’s recommended device-agnostic API, runs unmodified. A compiled model needs the Neuron backend named, once with torch.compiler.set_default_backend("neuron") , after which torch.compile(model) takes it, or per call as torch.compile(model, backend="neuron") . # Run a compiled training loop on Trainium torch.compiler.set_default_backend("neuron") # call once, before any compile model = MyModel().to('neuron') model.compile() # no backend argument optimizer = torch.optim.AdamW(model.parameters()) for batch in dataloader: optimizer.zero_grad() loss = criterion(model(batch), targets) loss.backward() optimizer.step() Capabilities TorchNeuron registers Trainium through PrivateUse1, PyTorch’s official extension point for custom hardware, so the standard APIs work as written: torch.nn , torch.tensor , torch.autocast , torch.distributed , and torch.profiler all operate on the neuron device. The two execution modes of PyTorch, eager and torch.compile work natively. High-level architecture diagram of TorchNeuron ( details ). Eager Eager runs operations immediately, so Python control flow, prints, and breakpoints all work with no compile or trace step. This is the mode for research and debugging. Eager mode dispatches operations individually. Executed directly, each operation compiles its own Neuron Executable File Format (NEFF) file, launches it on device, and moves its intermediates through HBM. We are also introducing adaptive eager mode (experimental), to reduce the launch cost and provide better opportunities to fuse operations to approach performance of compiled graph without a code change, while retaining the interactive experience. torch.compile torch.compile(model, backend="neuron") captures the model into a graph the Neuron compiler optimizes, and is the mode for stable, high-performance models moving to production. torch.compile can be used to optimize the whole or part of the model and works seamlessly with eager mode. FlexAttention is supported natively. TorchDynamo captures the model into an FX graph, AOTAutograd traces the backward graph ahead of time, the Neuron backend lowers the result, by way of Torch-MLIR, into StableHLO, and the Neuron compiler produces a NEFF file that executes on hardware. The Neuron inference stacks are built on the same torch.compile support in TorchNeuron. vLLM Neuron is a plugin to vLLM and serves LLMs; vLLM Omni Neuron is a plugin to vLLM Omni and serves diffusion models for image and video generation. Ops and dtypes Broad aten operator coverage for ops seen across text, multimodal and vision language, image and video generation and understanding, state space models, and world models. Models can use BF16, FP32, and FP16, with integer and boolean types, and torch.autocast and mixed precision code operate as written. TorchAO and FP8 are currently supported under private Beta and will be made generally available in the next release. Distributed APIs Distributed inference and training can use the standard PyTorch APIs and the high-level parallelism modules compose with your model as written: DDP, FSDP2, SimpleFSDP, tensor parallelism, and pipeline parallelism, along with common compositions like FSDP2 with tensor parallelism, ZeRO that are essential for training large models. DTensor and DeviceMesh are supported to handle sharding and the multi-dimensional process groups. The collective APIs (example: all_gather , reduce_scatter , all_reduce ) and the P2P APIs ( send , recv ) are supported. Collectives execute over Trainium’s dedicated collective cores intra-node and over the Elastic Fabric Adapter (EFA) inter-node, so the communication transport is native to the hardware with no change to PyTorch APIs. Streams and Events Torch streams and events are supported natively. torch_neuronx.Stream and torch_neuronx.Event map PyTorch’s stream semantics onto Trainium’s hardware queues. Within a stream, data movement and computation stay in order automatically. Separate streams carry no implicit dependency on one another, so across streams you make that ordering explicit by recording an event on the producing stream and waiting for it on the consuming stream before using the result. The device-agnostic torch.Stream , torch.Event , and torch.accelerator entry points reach the same implementation, so code written against them runs unmodified. An error surfaces at the next synchronization point rather than at the operation that caused it; TorchNeuron tracks the earliest exception and re-raises it there with its original context. In execution, 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; and when host collectives are disabled, a collective issued on a non-default stream keeps its place in that stream’s execution order but runs on the default hardware queue. Neuron Kernel Interface (NKI) The Neuron Kernel Interface (NKI) lets researchers and developers develop, optimize, and execute custom o