AWS opened up TorchNeuron today at PyTorch Conference North America in San Jose. It is a native PyTorch backend for Trainium2 and Trainium3, Apache-2.0, and it ships as a public beta in the Neuron 2.33.0 release. The pitch in the PyTorch blog post is that .to('neuron') is the only change you need.

I'm always a little suspicious of that sentence when it comes with new hardware. It isn't wrong. It just skips what operators care about: compile behavior, caches, shapes, and what happens when an op is missing.

What actually changed

The older torch-neuronx was built on PyTorch/XLA, so you wrote against torch_xla and xm.xla_device(). The new backend drops that. It registers Trainium through PrivateUse1, the same extension point PyTorch exposes for any out-of-tree device, so torch.nn, torch.autocast, torch.distributed and torch.profiler see a device called neuron. Code written against torch.accelerator picks it up without edits. You don't even need import torch_neuronx; the package registers a torch.backends entry point and autoloads on import torch.

The release notes call out the break directly: code written for the XLA-based library has to be updated. The docs also say the XLA lineage (versions before 2.14) stays available under the same package name. That's the first thing I'd pin in any requirements file, because torch-neuronx alone no longer tells you which stack you're getting.

Install surface

  • Packages: torch-neuronx 2.13.3.0 for PyTorch 2.13 and 2.14.3.0 for PyTorch 2.14, each with wheels for Python 3.12, 3.13 and 3.14 on Linux x86_64.
  • Install line: torch-neuronx==2.14.* plus neuronx-cc==2.* and nki from the Neuron pip repo, on top of the PyTorch CPU wheel. The device support all comes from the plugin.
  • System packages: the Neuron driver, runtime, collectives and tools, or a Neuron Base DLAMI that has them plus EFA.
  • OS: this is where the sources disagree. The docs release notes list Ubuntu 22.04 as end of support and 24.04 and 26.04 as supported. The GitHub v2.33.0 release says this release ships wheels for 26.04 only, with 22.04 and 24.04 wheels planned for the next release and build-from-source until then. If your base image is 24.04, check before you assume a wheel exists.

On a trn2.48xlarge, the quickstart shows neuron-ls listing 16 Neuron devices with 4 cores each and 96 GB per device, under a logical NeuronCore config of 2. PyTorch then reports 64 devices. Those values are Trn2-only; Trn3 needs its own math.

Eager mode is real, but every op is a compile

From the outside, eager behaves: control flow, prints and breakpoints all work. Underneath, each op is lowered to StableHLO, compiled to its own NEFF (Neuron Executable File Format) and queued asynchronously. A shape the compiler hasn't seen compiles on first use and is then served from cache.

Two details matter in practice:

  1. CPU fallback is quiet. An op without a Neuron implementation runs on the CPU with a warning, and the model keeps going. torch_neuronx.get_fallback_ops() lists what fell back. I'd check it at the end of every smoke test and fail CI if it isn't empty.
  2. Adaptive eager is opt-in and can change numerics. It fuses queued ops to cut launch overhead and is enabled with TORCH_NEURONX_ENABLE_CONCATENATION=1. The docs say fusion can change intermediate precision, and printing a tensor forces a sync that can break a fusion boundary. So a debug print can shift your numbers slightly. Worth knowing before you chase a loss mismatch.

torch.compile: no Inductor, static shapes

The compile path is TorchDynamo to an FX graph, AOTAutograd for the backward, a lowering to StableHLO through a vendored torch_mlir importer, then neuronx-cc produces the NEFF. TorchInductor isn't involved, so anything tuned for Inductor (Triton configs, max-autotune, CUDA graph modes) doesn't carry over.

The limits in this beta:

  • Only mode="default" and static shapes. You set dynamic=False, and each new input shape can trigger another compilation. The advice is to pad or bucket inputs. For training with fixed sequence length that's fine. For inference with variable prompts, bucketing is now your job unless a serving plugin does it for you.
  • Graph breaks work, lowering failures don't. Uncapturable code splits the graph into separately compiled segments, but a captured op the backend can't lower fails compilation by default.
  • FlexAttention is supported, and the docs suggest a mixed mode: run most of the model eager and compile hot regions like attention.

The compile cache is an operations problem

This is the section I'd hand to whoever owns the cluster images. Compiled NEFFs go into a content-addressed on-disk cache keyed on the lowered HLO bytes, the compiler configuration and the compiler version. A second cache holds the lowered StableHLO. The paths are TORCH_NEURONX_NEFF_CACHE_DIR and TORCH_NEURONX_HLO_CACHE_DIR, and the docs suggest pointing the NEFF cache at a shared mount such as FSx so ranks and instances can reuse artifacts.

Because the compiler version is part of the key, a neuronx-cc bump invalidates everything. If you let neuronx-cc==2.* float, your first job after an image rebuild pays the full compile cost on every rank. I'd pin the compiler, version the cache directory, and warm it in CI.

There's a tool for that warming: torch_neuronx.compiler.compile_only() or TORCH_NEURONX_COMPILE_ONLY=1 compiles graphs without running them. On a host with no Neuron device, you set NEURON_RT_FAKE_INSTANCE_TYPE to the target instance type. That lets a cheap CPU box pre-populate the cache before the Trainium nodes come up, which is the most useful operator feature in the release.

Distributed: standard APIs, NeuronCore as the unit

You initialize with init_process_group(backend="neuron"). The c10d backend covers all_reduce, all_gather, reduce_scatter, all_to_all, broadcast, reduce, gather, scatter and barrier, plus send, recv and batch_isend_irecv. On top of that: DDP, FSDP2 (fully_shard), SimpleFSDP, HSDP, tensor and sequence parallelism through DTensor and DeviceMesh, torch.distributed.pipelining, and experimental context parallelism for scaled dot-product attention. Collectives run on Trainium's collective cores inside a node and over EFA between nodes.

The thing that will trip up anyone porting GPU recipes is the unit of distribution. It's the NeuronCore, not the chip. The TorchTitan example in the blog trains Llama 3 8B on one trn2.48xlarge with NEURON_RT_NUM_CORES=64, torchrun --nproc_per_node 64, tensor parallel degree 4 and the rest sharded with SimpleFSDP. Your mesh shapes, ranks per node and memory per rank all come out of the NeuronCore count and the logical core setting, not the "8 GPUs per node" assumption baked into most configs.

Streams come with a caveat: non-default streams keep ordering and event semantics but multiplex onto a single hardware queue, so they don't buy real concurrency. Errors also surface at the next sync point, not at the op that caused them.

Dtypes and gaps

BF16 is the training baseline, with FP16 and FP32 supported. float64 runs as float32 with a warning, which matters for anyone with FP64 reductions in a loss or a metric. FP8 and MXFP8 require Trainium3. The release notes say this release doesn't include TorchAO integration or composite MXFP8 support, and the blog says TorchAO and FP8 are in private beta for the next release.

Not supported in this release: torch.sparse, torch.func, torch.nested and torch.distributions, and most of torch.fft falls back to the CPU. If you rely on torch.func for vmap or per-sample gradients, it won't run here yet.

How much is covered

AWS published its test methodology instead of a single compatibility number. It runs a sweep over 100,000 Hugging Face models, which it says cover 99.9% of real-world usage by download count, checking load, forward, backward, compile and which ops fell back to the CPU. It also runs PyTorch's upstream ATen, Dynamo, distributed and profiler suites. The August 2026 run produced about 154,000 test outcomes. Roughly 69,000 were skipped: about 30,000 by PyTorch's own eligibility checks (CUDA-only or CPU-only tests) and about 39,000 deliberately for Neuron. The largest group of those, about 22,000, are ops without a Neuron lowering yet.

On performance, the only concrete number in the announcement is from a customer. Deepgram's CTO says that in the benchmarks they evaluated, Trainium processes about 60% more input audio seconds per dollar than their existing GPU baselines, after a multi-month port with AWS's help. That's one team's workload and pricing, not a general result.

Why this matters beyond AWS

The same conference schedule has a session on TorchTPU, Google's native PyTorch backend for TPUs, which also lowers ATen to StableHLO and is described as heading toward a public GitHub repo. When two non-NVIDIA accelerators both ship as PrivateUse1 backends with eager mode, the cost of being off CUDA moves. It shifts from rewriting your training loop to managing compile caches, static shapes and op coverage. That's work a platform team can own once, instead of every model team paying for it.

What I'd do this week if I had Trainium capacity:

  1. Pin torch-neuronx to the exact 2.14.3.0 family and pin neuronx-cc, and confirm a wheel exists for your Ubuntu version.
  2. Run your model in eager mode and check get_fallback_ops() before touching torch.compile.
  3. Move the NEFF cache to shared storage and warm it with compile-only runs on a CPU host.
  4. Rebuild your device mesh from NeuronCore counts, not GPU counts.
  5. Bucket input shapes before measuring anything, or you'll be benchmarking the compiler.

It's a beta, and the release notes say so. But it's the first time Trainium has looked like an ordinary PyTorch device, and the parts that aren't ordinary are documented well enough to plan around.

Sources

Anpoo Sivanadi, Staff Software Engineer