AWS Open-Sources TorchNeuron: Trainium Becomes a Native PyTorch Device and Training Code Runs With .to('neuron'), but About 22,000 Upstream Tests Aren't Covered Yet

From "adapt it for Neuron" to "change the device name"

Running models on AWS's own Trainium chips used to mean adapting model code with Neuron-specific software and configuration, which is also how Hugging Face describes it in its statement. An AWS team wrote on the official PyTorch blog that TorchNeuron is open source under Apache 2.0 (repository aws-neuron/torch-neuronx) and in public beta from October 20, installable with pip install torch-neuronx from the AWS Neuron pip repository.

It uses PrivateUse1, the extension point PyTorch provides for third-party hardware, to register Trainium as a neuron device. In a training loop, .to('neuron') is the only change; torch.nn, torch.autocast, torch.distributed and torch.profiler work as written, and Transformers v5 loads models with device_map="neuron". Trainium2 and Trainium3 are supported.

Two modes and a kernel interface

Eager mode runs operations one at a time, so Python control flow and breakpoints work, which suits research and debugging; a new adaptive eager mode (experimental) tries to fuse operations and cut launch overhead without code changes. For production, torch.compile with the neuron backend goes through TorchDynamo, AOTAutograd and Torch-MLIR to StableHLO before the Neuron compiler takes over. For lower-level control, NKI lets you write custom operators, and the open-source NKI-lib ships common ones such as FlashAttention and MoE kernels. On the distributed side it supports DDP, FSDP2, tensor parallelism and pipeline parallelism, and vLLM has a Neuron plugin.

The gaps AWS lists itself

The post is upfront about coverage. An August run of PyTorch's upstream test suites produced about 154,000 outcomes, of which about 69,000 were skipped. About 30,000 of those are skipped by PyTorch's own eligibility checks (CUDA-only or CPU-only tests, for example); about 39,000 are deliberate Neuron skips, the largest group, about 22,000, being operators that don't have a native implementation yet. Sparse tensors, NestedTensor, forward-mode autograd and dynamic and symbolic shapes aren't supported in the first release, and TorchAO and FP8 are still in private beta.

Customer numbers so far come only through the vendor: 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. Teams with existing PyTorch code who want to compare GPU and Trainium costs can now try their own models directly; anyone relying on dynamic shapes or sparse operators should wait for later releases.