AWS 发布 TorchNeuron：面向 Trainium 的原生 PyTorch 后端

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 operators for both eager and torch.compile. Researchers can prototype ideas that require low level hardware control, including new model architectures. Developers can obtain finer performance control beyond Adaptive Eager and torch.compile by calling performance critical NKI kernels directly in model code to replace sequences of standard PyTorch operations (ATen ops).

TorchNeuron also provides Trainium optimized, high-performance operators implemented with NKI and released open source as NKI-lib, covering mainstream LLM and multimodal scenarios for both dense and MoE models, such as FlashAttention and MoE blockwise kernels, which developers can use to improve performance on Trainium. The following example defines and registers a custom NKI operator using the PyTorch recommended register_autograd style:

import torch
import nki
import nki.language as nl
from torch_neuronx import nki_op

# Declaring and implementing NKI kernels
# Note in NKI kernels, the output tensor should be defined as parameter to the kernel
@nki.jit
def sin_kernel(in_ptr0, out_ptr):
    input_tile = nl.load(in_ptr0[0:128])
    output_tile = nl.sin(input_tile)
    nl.store(out_ptr[0:128], value=output_tile)
    return out_ptr

# defining cos since it will be used in backward path of sin
@nki.jit
def cos_kernel(in_ptr0, out_ptr):
    input_tile = nl.load(in_ptr0[0:128])
    output_tile = nl.cos(input_tile)
    nl.store(out_ptr[0:128], value=output_tile)
    return out_ptr

# after this line, there is no NKI code, just native PyTorch, registering custom ops
@nki_op("mylib::pt_nki_sin", mutates_args={})
def pt_nki_sin(x: torch.Tensor) -> torch.Tensor:
    out = torch.empty_like(x)
    sin_kernel(x, out)
    return out

@nki_op("mylib::pt_nki_cos", mutates_args={})
def pt_nki_cos(x: torch.Tensor) -> torch.Tensor:
    out = torch.empty_like(x)
    cos_kernel(x, out)
    return out

def pt_nki_sin_backward(ctx, grad):
    x, = ctx.saved_tensors
    return grad * pt_nki_cos(x)  # cos is derivative of sin

def pt_sin_setup_context(ctx, inputs, output):
    x, = inputs
    ctx.save_for_backward(x)

# Register backward pass and context function used during autograd for nki_sin
pt_nki_sin.register_autograd(pt_nki_sin_backward, setup_context=pt_sin_setup_context)

if __name__ == "__main__":
    x = torch.randn(128, 1, device="neuron", requires_grad=True)
    # Forward pass uses sin_kernel
    y = pt_nki_sin(x)
    # Backward pass, autograd calls uses nki_sin_backward() -> cos_kernel()
    loss = y.sum()
    loss.backward()

Profiling with TorchNeuron

TorchNeuron integrates with the standard torch.profiler.profile API extended with a NeuronConfig object that controls Neuron-specific data collection. This allows users to capture and correlate traces from framework-level CPU operations (PyTorch dispatch, Python stacks) as well as Neuron runtime and Trainium activity (controlled by NeuronConfig) enabling end to end performance analysis.

Profiles captured can be visualized in the Neuron Explorer UI, where users can trace execution across framework, runtime, and device, jump from any event to the exact source line (source code viewer), navigate the workload through its call hierarchy (hierarchy viewer), and follow dependency chains to see what each stall is waiting on (dependency chain viewer).

Neuron Explorer also emits a Parquet output with per-op and per-engine metrics, so the same data behind the UI can be queried directly or consumed by agents for automated bottleneck analysis. For programmatic access, users can work with the built-in torch.profiler tools. For example, the key_averages table lets users drill down into specific ops and use the sort_by argument to surface the device-side cost of each op, which is useful for identifying hotspots.

Following code snippet shows an example (details):

import torch
from torch.profiler import profile, ProfilerActivity
from torch_neuronx.profiling import NeuronConfig, ProfileMode, NeuronProfiler

# Configure Neuron profiling and pass appropriate path as OUTPUT_DIR
nc = NeuronConfig(
   modes=[ProfileMode.DEVICE, ProfileMode.RUNTIME],
   profile_output_dir=OUTPUT_DIR,
   capture_enabled_for_nc='0',
   correlate_device_time=True,
)
# NeuronProfiler.export_trace() places the framework trace
# in the same directory as Neuron system traces
exporter = NeuronProfiler(nc)

device = torch.device('neuron')

x = torch.randn(4, 512, dtype=torch.bfloat16, device=device)
weight = torch.ones(512, dtype=torch.bfloat16, device=device)

def rms_norm_fn(t):
    x_sq = t * t
    mean_sq = x_sq.mean(dim=-1, keepdim=True)
    rms = torch.rsqrt(mean_sq + 1e-6)
    return t * rms * weight

# Profile your workload
with torch.profiler.profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.PrivateUse1],
    experimental_config=nc,
    record_shapes=False,
    on_trace_ready=exporter.export_trace,
    schedule=torch.profiler.schedule(wait=2, warmup=1, active=1),
) as prof:
    for step in range(4):
        with torch.no_grad():
            _ = rms_norm_fn(x)
        torch.neuron.synchronize()
        prof.step()

# Print the per-op summary, sorted by device time to surface hotspots
print(prof.key_averages().table(sort_by="self_privateuse1_time_total", row_limit=4))

For additional capabilities and features refer to Neuron Explorer documentation.

Customers and partners building on TorchNeuron

TorchNeuron is already in use across startups, enterprises, ML industry partners, and academic institutions: spanning autonomous vehicles, robotics, video generation, speech, telecommunications, healthcare, and geospatial modeling. Below is a snapshot of what teams are building today.

“Cartesia is building intelligence that learns from and interacts with the world the way humans do, starting with conversation. Our pioneering work on state-space architectures like Mamba powers Sonic (text-to-speech) and Ink (speech-to-text), the voice AI models behind conversational experiences that feel truly human. Real conversation demands three things at once: quality (naturalness and accuracy), efficiency (low latency at scale), and steerability. Most solutions deliver one or two. We designed our research and products from first principles to deliver all three, which means we need infrastructure that lets us squeeze maximum performance out of every layer without rewriting our stack. With PyTorch native support on AWS Trainium, we ran our standard inference pipeline with minimal changes to our PyTorch code. Small, fast ops map efficiently through the native path, so we didn’t have to hand-write kernels for them. That lets our team concentrate on writing high-performance NKI kernels for the layers that are called repeatedly and drive end-to-end latency. For real-time voice, every millisecond of network distance matters. Running on Trainium gives helps to put Sonic and Ink close to our customers globally.”

Arjun Desai, Cofounder, Cartesia

“Voice is becoming a first-class modality for AI, and Deepgram processes quadrillions of audio tokens a year for more than 200,000 developers and 1,400 organizations. To meet that demand, we need to scale efficiently while continuing to ship better models. Over a multi-month engagement, the AWS Generative AI Innovation Center helped us port our model to PyTorch Native on AWS Trainium. In the benchmarks we evaluated, Trainium processes about 60% more input audio seconds per dollar than our existing GPU baselines, which helps us bring better models to customers at lower prices. We look forward to supporting future AWS AI chips in our solution.”

Adam Sypniewski, CTO, Deepgram

“Hugging Face is the leading open platform for AI builders, providing free access to over 6 million open models, datasets and applications through PyTorch-powered open-source libraries including Transformers and Diffusers. Until now, taking advantage of AWS Trainium required adapting Hugging Face models with Neuron-specific software and configuration; TorchNeuron makes Trainium a native backend for Transformers and Diffusers, so developers can use familiar model definitions through a fully open-source stack – a practical expression of our commitment to silicon neutrality and compatibility. It’s now easier than ever to deploy state-of-the-art open LLMs and the latest open image and video generation models on AWS Trainium, with TRL enabling fine-tuning and reinforcement learning and Neuron kernels delivering optimized performance automatically through our libraries.”

Jeff Boudier, VP Product, Hugging Face

Beyond these customers, we have been working with a number of partners who are building with TorchNeuron to enhance the Trainium ecosystem. These partners include Anyscale, Decart, Kernelize, Loka, Luminal, Makora, Modular, Nota AI, Radix, Wafer, and Yasp.

How TorchNeuron Enables broader compatibility

One of our goals with native PyTorch support is to enable developers to leverage the performance and cost benefits of Trainium while continuing to use the models, libraries, and workflows they already know. This helps preserve their investment in PyTorch, including their model code, inference and training infrastructure, debugging practices, and the community projects they build upon.

The cost of running AI workloads includes the engineering time needed to get them working, optimize their performance, and maintain them as models and libraries evolve. We want developers to spend more time experimenting with architectures and improving model quality, and less time adapting their software to the underlying hardware. They should be able to start with a familiar model or training recipe, adapt it to their workload, and run it on Trainium with minimal changes specific to Neuron.

TorchNeuron provides the foundation for this experience through compatibility with the broader PyTorch. Developers can use libraries such as TorchTitan, which brings together PyTorch’s distributed and compilation capabilities for large-scale training, and Hugging Face Transformers, which provides access to community model implementations, pretrained checkpoints, and APIs for adapting models to their applications.

TorchTitan Example

TorchTitan is PyTorch’s native library for rapid experimentation and large-scale training. Its modular design makes it easier to customize training, test new ideas, and scale up experiments. TorchTitan trains on Trainium today. We continuously upstream into mainline TorchTitan, and torchtitan-neuron gives you those improvements before they merge, converging on mainline as each one lands. The Trainium-specific artifacts also ship in torchtitan-neuron and register through TorchTitan’s own extension points; those include NKI kernels, Trainium-tuned model configurations and training recipes, and profiling integrations for Neuron Explorer.

# Install from source

git clone https://github.com/aws-neuron/torchtitan-neuron
cd torchtitan
pip install -r requirements.txt

# Download the llama3-8B tokenizer

python scripts/download_hf_assets.py --repo_id meta-llama/Llama-3.1-8B --assets tokenizer --hf_token=...

# Train Llama3-8B with Simple FSDP + torch.compile on a single trn2.48xlarge instance:

export NEURON_RT_NUM_CORES=64
torchrun --nnodes 1 --nproc_per_node $NEURON_RT_NUM_CORES \
--local-ranks-filter 0 --role rank --tee 3 \
--rdzv_backend c10d --rdzv_endpoint "localhost:0" \
-m torchtitan.train --module torchtitan_neuron.experiments.graph_trainer.llama3 --config simple_fsdp_llama3_8b \
--parallelism.tensor-parallel-degree 4 \
--parallelism.data-parallel-shard-degree -1 \
--training.global-batch-size 16 \
--training.local-batch-size 1 \
--compile.enable --compile.backend neuron

Following is example loss curve of training Llama 3 8B on 8 x trn2.48xlarge with Tensor Parallel (TP) and Hybrid Sharded Data Parallel (HSDP) with TorchTitan

Hugging Face Transformers v5

Transformers v5 is built for interoperability: one model implementation that runs on any platform and any hardware, and that downstream training and serving stacks consume directly.

Trainium is a supported device upstream in the Hugging Face libraries, including Transformers, Diffusers, trl and kernels. Transformers v5 enables training and inference of text and audio models, Diffusers supports image and video generation, trl enables reinforcement learning of large language models, and NKI kernels are available through the Hugging Face kernels library. AWS and Hugging Face added the support to each library, so existing Hugging Face workflows run on Trainium with no device-specific code.

Training. Transformers v5 strengthens full pre-training and fine-tuning readiness, with better initialization and compatibility with the parallelism and tooling ecosystem — TorchTitan for scale, Unsloth and TRL for post-training.
Inference. Transformers v5 standardizes attention and kernels across model implementations.
Diffusion. Tensor parallelism in Diffusers is merged upstream for Neuron, so image and video generation run through the diffusers APIs unchanged.
Compatibility. Transformers v5 is a native backend of runtime libraries such as TorchTitan, vLLM and SGLang, extending their model breadth to most of the 3 million public models available on Hugging Face. We are working on enabling and validating this integration on Neuron, bringing this expanded model support to Trainium instances.

Loading a model onto Trainium

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    # Neuron hardware is automatically detected
    device_map="neuron"
)

TorchNeuron registers Trainium through PyTorch’s PrivateUse1 extension point, which puts Neuron in PyTorch’s device enumeration. device_map=”auto” finds it with no configuration, and device_map=”neuron” names it explicitly. Either way the weights load onto the device for training or inference, and the rest of the script is unchanged.

Testing methodology and results

Our testing methodology included verifying TorchNeuron on broad model coverage across HuggingFace and TorchTitan models, and PyTorch upstream test suite across both eager and torch.compile mode.

Model coverage

Blind Sweep (Broad testing at scale). We used the Hugging Face model universe to perform broad sweep across open source models. We continuously run 100k Hugging Face models that cover 99.9% of real-world usage (based on the download count). We ran each model on Trainium in a reduced but representative configuration and record what happens: does it load, does the forward and backward pass run, does it compile, and which operators executed on the device versus fell back to CPU.

The following graph shows the result of 100k model run across eager and torch.compile mode as of September, 2026. We are working on closing the remaining gaps, targeting closing most, if not all, in future releases.

PyTorch Upstream tests

We run the PyTorch upstream test suites: ATen, Dynamo, distributed, and profiler, and report the pass rate against the enabled tests. The August 2026 run produced about 154K test outcomes, of which roughly 69K were skipped with a valid reason: About 30K of those are skipped by PyTorch own eligibility checks (for example CUDA-only and CPU-only tests) and are not Neuron-specific. The remaining ~39K are deliberate Neuron skips, in three groups

Operators without a native Neuron lowering at this point (~22K, the largest group; ties to operator coverage above that we plan to close in sequential releases) – corner cases for elementwise, reduction, foreach, advanced-indexing, and advanced-linalg operators, and fused LSTM/GRU cells.
CUDA-specific tests (i.e. the upstream eligibility check did not flag) – CUDA graphs, NCCL-specific collectives, CUDA-specific Dynamo paths tests are skipped with eligibility check.
Features TorchNeuron explicitly chose not to support in the first release as customers didn’t request them: composite-operator compliance, sparse tensors and sparse ops, NestedTensor, forward-mode autograd (JVP), arbitrary attention masks in native multi-head attention (no-mask and causal are supported), per-channel quantization, negative/conjugate view bits, torch.signal windows, bilinear/bicubic up sample, and dynamic and symbolic shapes.

As part of this process, we also contributed back to PyTorch by raising several RFCs and PRs to improve test quality and removing hard coded CUDA assumptions (193439, 193441, 193442, see Appendix A for complete PR list).

How to get started

Starting Oct. 20, 2026, TorchNeuron will be available as a public beta. Researchers and developers will be able to install with pip install torch-neuronx from the AWS Neuron pip repo details.

Resources

The following resources will go live on Oct. 20:

TorchNeuron: https://github.com/aws-neuron/torch-neuronx
AWS Trainium: https://aws.amazon.com/machine-learning/trainium/
AWS Neuron SDK: https://aws.amazon.com/machine-learning/neuron/
Neuron Documentation: https://awsdocs-neuron.readthedocs-hosted.com/

Quick Reference

For further guidance on getting started, execution modes, distributed training, and development roadmap, see our Frequently Asked Questions (FAQ) section in the Neuron documentation.

Join us at PyTorch Conference North America

Open source TorchNeuron under Apache 2.0 on GitHub
Release pip-installable packages for immediate use
Present technical deep-dive sessions on native PyTorch on Trainium
Host live demos

Acknowledgments

We would like to thank the PyTorch community and the PyTorch team at Meta for their work on PyTorch’s extensibility mechanisms, particularly the PrivateUse1 device backend and torch.accelerator APIs that made this integration possible. We also thank the PyTorch Foundation for their support in CI/CD infrastructure.

AWS is committed to supporting the PyTorch ecosystem by providing Trainium resources for PyTorch CI/CD infrastructure and contributing to open source PyTorch projects.

Special thanks to Anthropic, Anyscale, Amazon AGI, Amazon Rufus, ByteDance, Cartesia, Clay Foundation, Databricks, Decart, Deepgram, Hugging Face, Karakuri, Lieber Institute for Brain Development, Meta, MIT, Neura Robotics, NinjaTech AI, Nota AI, Odyssey, PDT Partners, poolside, Qube Research and Technologies, RadixArk, Ricoh, Rivian, Salesforce, SemiAnalysis, UC Berkeley, Stanford, CMU, UCLA, Zyphra, and many others for their feedback in shaping TorchNeuron’s design and capabilities.

Contributors: Apoorv Gupta, Arash Jamalian, Rui Silva, Peter Sam Raj, Kunind Sahu, Kamran Khan, Shreya Srivastava, Charan Shettyhalli Guruswamy, Alex Hoffman, Abigail Fernandes, Sam Jin

Appendix: Complete List of PRs

The post Introducing TorchNeuron: Native PyTorch backend for AWS Trainium appeared first on PyTorch.