Amazon's TorchNeuron is a new way to run PyTorch on AWS Trainium chips. It exposes Trainium as a standard PyTorch device called neuron, so the documented first step is to change .to('cuda') to .to('neuron') instead of porting a model to a separate graph-compilation API. This explainer covers what TorchNeuron is, which hardware it targets, what the beta does and does not do, how to install it as AWS documents it today, and how it compares with the older torch-neuronx XLA path and with CUDA. For the argument about what this means for Nvidia, see our opinion piece, Amazon's TorchNeuron Attacks the Right Part of Nvidia's Moat, but Not the Hard Part. This post sticks to the mechanics.
A note on timing. As of Monday, October 19, 2026, AWS sources differ on access. The AWS team's launch post on the PyTorch blog says TorchNeuron is "going live as public Beta on Oct. 20" and that "Starting Oct. 20, 2026, TorchNeuron will be available as a public beta." The Neuron documentation page for TorchNeuron still says "TorchNeuron is currently only available as part of a closed Beta program." The GitHub repository the launch post names, github.com/aws-neuron/torch-neuronx, returned HTTP 404 when we checked at about 10:05 AM PT today. The launch post lists it among resources that "will go live on Oct. 20."
What TorchNeuron is
The Neuron docs define it as "an open-source native PyTorch backend for AWS Trainium that integrates through PyTorch's standard PrivateUse1 device backend mechanism." The launch post calls PrivateUse1 "PyTorch's official extension point for custom hardware," and says that because TorchNeuron registers through it, torch.nn, torch.tensor, torch.autocast, torch.distributed, and torch.profiler "all operate on the neuron device."
The launch post says "The two execution modes of PyTorch, eager and torch.compile work natively":
- Eager mode. Operations "are dispatched and execute immediately upon invocation," according to the docs. PyTorch's dispatcher routes tensor operations to the Neuron backend, which provides implementations of ATen operators. The launch post calls eager "the mode for research and debugging," where "Python control flow, prints, and breakpoints all work with no compile or trace step."
- torch.compile. TorchNeuron implements a custom backend for TorchDynamo. Per the docs, TorchDynamo captures an FX graph, AOT Autograd generates forward and backward graphs, the Neuron backend lowers them to Neuron IR, and the Neuron Compiler generates Trainium instructions. The launch post describes the lowering as going "by way of Torch-MLIR, into StableHLO," with the compiler producing a NEFF file. The docs say the backend provides
torch.compilesupport "without using Inductor" for this release.
On top of plain eager mode, AWS adds Adaptive Eager Execution. The docs say it analyzes sequences of queued operators and fuses them, "improving performance while maintaining the same execution order, numerical accuracy, and determinism as non-fused execution." The launch post labels adaptive eager mode "(experimental)." Per the docs, when torch.use_deterministic_algorithms() or torch.set_deterministic_debug_mode() is called, "Adaptive Eager Execution optimizations are disabled."
For hand-tuned code there is the Neuron Kernel Interface (NKI). The docs say NKI "uses similar definition and registration patterns as Triton," with kernels defined using the @nki.jit decorator and registered as PyTorch operators with @nki_op. Training models that include NKI kernels "requires a backward version of the custom op, implemented using the register_autograd() function." AWS also publishes prebuilt kernels in NKI-lib, which the launch post says covers mainstream LLM and multimodal scenarios, "such as FlashAttention and MoE blockwise kernels." For how Triton compares with CUDA C++, CUTLASS and torch.compile on GPUs, see our guide to custom LLM kernel optimization.
TorchNeuron has existed in some form since late 2025. The About PyTorch on AWS Neuron page says it "was announced in December 2025 at AWS re:Invent and shipped its initial version as part of Neuron release 2.27.0." The same page says it "retains the same Python package name as its predecessor (torch-neuronx)" while being "an entirely new native PyTorch backend." That shared name matters during setup, covered below.
Which chips and instances it supports
TorchNeuron targets Trainium2 and Trainium3. The Neuron docs page is marked relevant for Trn2 and Trn3, and the About page lists "Target Platforms: AWS Trainium (Trn2, Trn3 instances)." Amazon's announcement says the same: "TorchNeuron supports Trainium2 and Trainium3."
The Neuron architecture docs describe these configurations:
- Trn2. The Trn2 architecture page covers the trn2.48xlarge and trn2u.48xlarge instances, which "are powered by 16 Trainium2 chips," and the Trn2 UltraServer, which "comprises four trn2u.48xlarge instances" for a total of 64 Trainium2 chips. The page lists 1,536 GiB of device memory per instance.
- Trn3. The Trn3 architecture page says "Trn3 is available in two UltraServer scale-up configurations: Gen1 with 64 Trainium3 chips per UltraServer, and Gen2 with 144 chips per UltraServer."
Not on the TorchNeuron list: Inferentia (Inf1, Inf2) and first-generation Trainium (Trn1, Trn1n). The About page says torch-neuron is for Inf1, and torch-neuronx, the XLA-based package, targets "AWS Inferentia2 (Inf2 instances) and AWS Trainium (Trn1, Trn1n, Trn2, Trn3 instances)." If your fleet is Inf2 or Trn1, the docs point you to the XLA path, not TorchNeuron.
One sizing detail changes how you configure parallelism. "On Trainium, the unit of distribution is the NeuronCore," per the docs FAQ, and its TorchTitan example says to "set NGPU=64 for 16 Trainium2 chips (4 NeuronCores per chip)."
What works in the beta
Per the Neuron docs and the launch post:
- Distributed APIs. The docs list Fully Sharded Data Parallel (FSDP), Distributed Tensor (DTensor), Distributed Data Parallel (DDP), and Tensor Parallelism (TP). "For eager mode, TorchNeuron supports all three: FSDPv1, FSDPv2, and SimpleFSDP. For torch.compile, TorchNeuron follows the PyTorch community recommendation and supports SimpleFSDP." The launch post says collectives "execute over Trainium's dedicated collective cores intra-node and over the Elastic Fabric Adapter (EFA) inter-node."
- Data types. "Models can use BF16, FP32, and FP16, with integer and boolean types," per the launch post. For mixed precision you can pass a policy to FSDP or use
torch.autocast(device_type="neuron"). The docs say "Both are supported." - Activation checkpointing. "Yes. TorchNeuron supports activation checkpointing."
- Graph breaks under
torch.compile: "Yes, the Neuron TorchDynamo backend supports graph breaks." - MoE ops. The docs say "torch native MoE ops will be supported from first release, including torch._scaled_grouped_mm, torch._grouped_mm, and MoE Dispatch/Combine operations," and "First TorchNeuron release already comes with GPT-OSS support that covers all this."
- FlexAttention "is supported natively," per the launch post.
- Streams and events. "torch_neuronx.Stream and torch_neuronx.Event map PyTorch's stream semantics onto Trainium's hardware queues."
- Profiling. TorchNeuron "integrates with the standard torch.profiler.profile API extended with a NeuronConfig object," and profiles open in the Neuron Explorer UI.
- Ecosystem. The launch post says "Trainium is a supported device upstream in the Hugging Face libraries, including Transformers, Diffusers, trl and kernels," and "TorchTitan trains on Trainium today." For serving, "vLLM Neuron is a plugin to vLLM and serves LLMs," and the Neuron inference stacks "are built on the same torch.compile support in TorchNeuron."
AWS also published test coverage. The launch post says "We continuously run 100k Hugging Face models that cover 99.9% of real-world usage (based on the download count)," recording whether each model loads, whether forward and backward passes run, whether it compiles, and "which operators executed on the device versus fell back to CPU." On PyTorch's own upstream suites, "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's own eligibility checks, and "The remaining ~39K are deliberate Neuron skips."
What does not work yet
The beta has real gaps. From the Neuron docs FAQ:
- No CUDA graphs equivalent. Asked whether TorchNeuron has equivalents to CUDA graphs and CUDA graph trees, the docs answer: "Not in the initial release. We are considering equivalent constructs for future releases."
- Compilation needs Trainium hardware. "The initial release requires compilation on Trainium2 or Trainium3 (Trn2 or Trn3) instances. Future releases will support compilation on non-Trainium instances."
- Some MoE collectives later.
torch.all_to_all_vdev_2dandtorch.all_to_all_vdev_2d_offset"will be supported in future releases." - PyTorch Foundation libraries. For torchcomms, monarch, torchforge, and torchao, AWS says "We are actively evaluating support for these libraries now. Our goal is to support all of them over the next couple of quarters."
- Out of tree. "Currently, TorchNeuron is an out-of-tree backend." The docs say "Our ultimate goal is to move in-tree."
From the launch post:
- TorchAO and FP8 "are currently supported under private Beta and will be made generally available in the next release."
- Operator gaps. The largest group of Neuron test skips, about 22K, is "Operators without a native Neuron lowering at this point," covering "corner cases for elementwise, reduction, foreach, advanced-indexing, and advanced-linalg operators, and fused LSTM/GRU cells."
- Features left out of 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."
- Stream behavior. "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." Also, "An error surfaces at the next synchronization point rather than at the operation that caused it."
The missing support for dynamic and symbolic shapes is the item to check first if your models depend on them.
Pipeline parallelism: the sources disagree. The Neuron docs say "Support for additional parallelism strategies such as Pipeline Parallelism (PP) will be available soon." The launch post lists "pipeline parallelism" among the parallelism modules that "compose with your model as written." Test it before you depend on it.
How to try it
The Neuron setup pages we checked on October 19 do not yet have a separate TorchNeuron install section. The steps below use only what AWS documents today, and say which page each step comes from.
1. Get access. Until the public beta opens, the docs say: "If you would like to participate, contact your AWS Neuron support representative." From October 20, the launch post says "Researchers and developers will be able to install with pip install torch-neuronx from the AWS Neuron pip repo," and links to the Install PyTorch for Neuron page for details.
2. Launch a Trn2 or Trn3 instance. The manual install page lists Ubuntu 24.04, Ubuntu 22.04, or Amazon Linux 2023, and Python 3.11, 3.12 or 3.13. It also says: "When launching a Trn1/Trn2/Trn3 instance, you must adjust your primary EBS volume size to a minimum of 512GB." The same page says "Ubuntu 22.04 has reached end-of-support on Neuron," "New deployments should use Ubuntu 24.04," and "Currently, PyTorch 2.9 is not available on Amazon Linux 2023." The setup index also offers the Neuron multi-framework Deep Learning AMI and Deep Learning Containers as alternatives to a manual install.
3. Install the Neuron driver, runtime, and tools. The manual page gives these commands for Ubuntu 24.04 (labeled "Neuron 2.32 versions"):
# Configure Linux for Neuron repository updates
. /etc/os-release
sudo tee /etc/apt/sources.list.d/neuron.list > /dev/null <<EOF
deb https://apt.repos.neuron.amazonaws.com ${VERSION_CODENAME} main
EOF
wget -qO - https://apt.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB | sudo apt-key add -
# Update OS packages
sudo apt-get update -y
# Install OS headers
sudo apt-get install linux-headers-$(uname -r) -y
# Install git
sudo apt-get install git -y
# Install Neuron Driver (Neuron 2.32 versions)
sudo apt-get install aws-neuronx-dkms=2.30.2.0* -y
# Install Neuron Runtime
sudo apt-get install aws-neuronx-collectives=2.34.10.0* -y
sudo apt-get install aws-neuronx-runtime-lib=2.34.10.0* -y
# Install Neuron Tools
sudo apt-get install aws-neuronx-tools=2.32.28.0* -y
# Add PATH
export PATH=/opt/aws/neuron/bin:$PATHThe page adds an EFA install step for Trn1, Trn1n, Trn2, and Trn3, marked "only required for multi-instance training."
4. Point pip at the Neuron repository and install. The manual page creates a virtual environment and sets the extra index like this:
python3.12 -m venv aws_neuron_venv_pytorch
source aws_neuron_venv_pytorch/bin/activate
python -m pip install -U pip
python -m pip config set global.extra-index-url https://pip.repos.neuron.amazonaws.comThen, per the launch post, from October 20:
pip install torch-neuronxWatch the version pin. The manual install page pins torch-neuronx==2.9.0.2.15.32035 under the comment "torch-neuronx pinned to Neuron 2.31 as XLA support ended with that release," and says "torch-neuronx (PyTorch/XLA) is not included in Neuron 2.32.0." Because both packages share the name torch-neuronx, that pinned build is the PyTorch/XLA package, not TorchNeuron.
5. Verify the devices. The manual page runs neuron-ls, which prints a table of Neuron devices with their cores and memory.
6. Change your code. Per the docs FAQ, for eager mode:
- "Device placement: Change .to('cuda') to .to('neuron')"
- "If your code uses torch.accelerator, no changes are needed (automatic device detection)"
- "Mixed precision: Use standard torch.autocast(device_type="neuron") API"
- Set sharding by NeuronCore count where a strategy needs it. "For some parallelism strategies like Tensor Parallelism, you need to specify how many NeuronCores are used for sharding. For other strategies like FSDP, no configuration changes are needed."
For torch.compile, the docs add two changes: "Specify backend="neuron"" and "Remove CUDA-specific parameters like mode="max-autotune-no-cudagraphs"". The launch post shows the minimal training loop:
# 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()To compile, the launch post calls torch.compiler.set_default_backend("neuron") once, "before any compile," then model.compile() with no backend argument. The docs say "the torch_neuronx import is only needed when using NKI kernels," because TorchNeuron uses the backend autoloading that PyTorch introduced in PyTorch 2.9.
With Hugging Face Transformers, the launch post loads Qwen/Qwen3-0.6B with AutoModelForCausalLM.from_pretrained(..., device_map="neuron") and says "device_map="auto" finds it with no configuration." For TorchTitan in eager mode, the docs say "Zero code changes required." AWS keeps Trainium-specific recipes and NKI kernels in torchtitan-neuron. The launch post's example will "Train Llama3-8B with Simple FSDP + torch.compile on a single trn2.48xlarge instance" with NEURON_RT_NUM_CORES=64.
TorchNeuron vs the Neuron SDK XLA path (torch-neuronx)
Before TorchNeuron, PyTorch training on Trainium meant torch-neuronx built on PyTorch/XLA. Per the About page:
- Device type. The XLA path integrates as a "PyTorch/XLA device backend (xla device type)." TorchNeuron is a "Native PyTorch backend via PrivateUse1 mechanism (neuron device type)."
- Execution model. The XLA path uses "XLA-based compilation with support for dynamic shapes and control flow." TorchNeuron "provides eager mode execution, torch.compile support, and standard PyTorch distributed APIs without requiring XLA."
- Hardware. XLA path: Inf2, Trn1, Trn1n, Trn2, Trn3. TorchNeuron: Trn2 and Trn3.
- Parallelism. The XLA path has "Native support for data parallelism, tensor parallelism, pipeline parallelism, sequence parallelism, and Zero Redundancy Optimizer (ZeRO)," and works with libraries such as NeuronX Distributed (NxD). TorchNeuron uses standard PyTorch APIs: FSDP, DTensor, DDP, and TP.
- PyTorch versions. The page says to choose the XLA path when "Using PyTorch versions 2.5 through 2.9 (XLA-based implementation)," and that "PyTorch 2.9 is the last version of torch-neuronx based on PyTorch/XLA." It says to choose TorchNeuron when "Training models on AWS Trainium (Trn2, Trn3) instances with PyTorch 2.10 or later."
AWS is moving users off the XLA path. The PyTorch setup page says "The torch-neuronx library is no longer included on Neuron DLAMIs and DLCs starting with Neuron release v2.32.0," and "If you are using torch-neuronx for training your models, pin your environment to Neuron v2.31.0." The About page's migration note says "Starting with PyTorch 2.10 support (planned for a future Neuron release), AWS Neuron will transition from PyTorch/XLA to native PyTorch support via TorchNeuron. Users on PyTorch 2.9 or earlier will need to update their scripts when upgrading to PyTorch 2.10 or later."
The trade-off, as the docs describe it: the XLA path has dynamic shape support and covers Inferentia2 and Trn1, and TorchNeuron's first release has neither. TorchNeuron adds eager-mode debugging and standard PyTorch distributed APIs, and the About page says it "eliminates the need for XLA-specific APIs."
TorchNeuron vs CUDA
TorchNeuron gives Trainium the same kind of PyTorch device interface that CUDA GPUs already have. It does not make the two stacks identical. Based on AWS's own descriptions:
- Same PyTorch surface. The device string changes,
torch.acceleratorcode needs no changes, and autocast uses "automatic datatype conversion following PyTorch CUDA conventions," per the docs. - Different compiler path. The TorchNeuron
torch.compilebackend lowers graphs to Neuron IR for the Neuron Compiler, without using Inductor in this release. - No CUDA graphs. CUDA-specific compile parameters must be removed, and there is no CUDA graph equivalent in the initial release. For what CUDA graphs do in GPU serving, see our explainer on CUDA graph capture in production LLM serving.
- Different parallelism unit. You size jobs by NeuronCores. The docs note the number of NeuronCores "can be different from the number of GPUs."
- Kernels. NKI follows "similar definition and registration patterns as Triton." The launch post quotes Cartesia saying "Small, fast ops map efficiently through the native path, so we didn't have to hand-write kernels for them," while its team concentrates "on writing high-performance NKI kernels for the layers that are called repeatedly."
- Tests. AWS skips "CUDA graphs, NCCL-specific collectives, CUDA-specific Dynamo paths tests," per the launch post.
On cost, AWS's evidence so far is customer-reported. In the launch post, Deepgram CTO Adam Sypniewski says "In the benchmarks we evaluated, Trainium processes about 60% more input audio seconds per dollar than our existing GPU baselines," after "a multi-month engagement" in which the AWS Generative AI Innovation Center helped port the model. That is one company's workload. Amazon's announcement describes a demonstration where engineers "ran the same AI training script on a conventional chip and on Trainium" and "Results showed that the AI model learned the same things at the same rate." The announcement does not publish the underlying numbers.
Who should try it now
If you run PyTorch on Trn2 or Trn3, TorchNeuron is the path AWS says it is moving to, and the install page names Neuron 2.31 as the last release that ships the XLA package. A reasonable first test is a model with static shapes that uses standard distributed APIs: run it in eager mode, then compile the parts that need speed. If you need dynamic shapes, Inferentia2, or Trn1, the docs point you to the XLA path for now. Sparse tensors are not in the first release either. Treat everything here as beta behavior: AWS lists several of these gaps for future releases.
Sources
- Native PyTorch for AWS Trainium (AWS Neuron documentation)
- About PyTorch on AWS Neuron (AWS Neuron documentation)
- Introducing TorchNeuron: Native PyTorch backend for AWS Trainium (AWS team, PyTorch blog)
- Amazon opens its AWS Trainium AI chips to every PyTorch developer, no rewrite needed (About Amazon)
- Install PyTorch via manual installation (AWS Neuron documentation)
- Install PyTorch for Neuron (AWS Neuron documentation)
- Install AWS Neuron SDK (AWS Neuron documentation)
- Amazon EC2 Trn2 Architecture (AWS Neuron documentation)
- Amazon EC2 Trn3 Architecture (AWS Neuron documentation)
- NKI Library (aws-neuron on GitHub)
- torchtitan-neuron (aws-neuron on GitHub)
- Opinion: Amazon's TorchNeuron Attacks the Right Part of Nvidia's Moat, but Not the Hard Part (llms.blog)
