AWS opens PyTorch to Trainium chips with one-line model port

Summary

• AWS released TorchNeuron, an open-source backend running standard PyTorch code directly on Trainium silicon.
• Moving a model requires changing a single line: model = MyModel().to('neuron'), with no Neuron-specific rewrite.
• It gives teams a practical in-AWS alternative to Nvidia GPUs without abandoning existing PyTorch workflows.

The latest

Standard PyTorch code can now run on Amazon's Trainium chips as an ordinary device, after AWS and the PyTorch team published TorchNeuron on 19 October 2026. The backend is built on PyTorch's accelerator-neutral extension points, the PrivateUse1 device interface and the torch.accelerator APIs, and supports Trainium2 and Trainium3. A public preview opens 20 October.

Details

• The port: Moving an existing model takes one line, model = MyModel().to('neuron'), according to the PyTorch blog, while compiled models switch with torch.compiler.set_default_backend("neuron"). The stated design goal is that no Neuron-specific rewrite is required, with Trainium treated as a regular PyTorch device.
• Execution modes: AWS said eager mode covers research and debugging, with an experimental adaptive eager mode that lowers launch cost. Production runs through torch.compile, following a path from TorchDynamo to AOTAutograd to Torch-MLIR (StableHLO) and then the Neuron compiler. FlexAttention is supported.
• Custom kernels: The NKI kernel interface lets developers write custom kernels, and AWS open-sourced NKI-lib, a library including FlashAttention and mixture-of-experts kernels. Cartesia co-founder Arjun Desai said his team wrote NKI kernels only for recurrent layers, leaving the rest of the inference pipeline largely unchanged.
• Coverage: torch.nn, torch.tensor, torch.autocast, torch.distributed and torch.profiler all operate on the neuron device, per the blog. Supported data types are BF16, FP32, FP16, integers and booleans. TorchAO and FP8 remain in private preview, with general availability promised for the next release.
• Distributed training: The backend supports DDP, FSDP2, SimpleFSDP, tensor parallelism, pipeline parallelism, DTensor and DeviceMesh. Intra-node collectives run on Trainium collective kernels; inter-node traffic uses Elastic Fabric Adapter. Profiling traces link framework, runtime and device activity to source lines in Neuron Explorer, which exports Parquet data.
• The numbers: AWS said it continuously runs roughly 100,000 Hugging Face models, which it says represent 99.9% of real-world usage by download count, each in a miniature configuration testing load, forward and backward passes, compilation and CPU fallback. September 2026 results were published as a chart without a stated pass rate.
• Test gaps: On native PyTorch tests from August 2026, about 154,000 results included roughly 69,000 skips. Around 30,000 were PyTorch's own eligibility skips; about 39,000 were Neuron-specific, including some 22,000 for operators lacking a native lowering. Sparse tensors, NestedTensor, forward-mode autodiff, dynamic shapes and arbitrary attention masks were excluded from the first release.
• Ecosystem: TorchTitan trains on Trainium via torchtitan-neuron, with examples covering Llama 3 8B on a single trn2.48xlarge instance and an eight-instance run using tensor parallelism and HSDP. Hugging Face Transformers v5 and Diffusers accept device_map=neuron. Inference runs through vLLM Neuron and vLLM Omni Neuron.
• Customer claim: Deepgram chief technology officer Adam Sypniewski said Trainium processes roughly 60% more input audio seconds per dollar than the company's existing GPU baselines on the benchmarks evaluated, after a migration that took months with the AWS Generative AI Innovation Center. Deepgram says it serves more than 200,000 developers and 1,400 organizations.
• Licensing: TorchNeuron ships under Apache 2.0 at github.com/aws-neuron/torch-neuronx, installed via pip install torch-neuronx. The public preview begins 20 October 2026, timed to the PyTorch Conference North America in San Jose on 20-21 October. AWS listed roughly 30 upstream merge requests to PyTorch and TorchTitan.

Background

Trainium is Amazon's in-house AI training silicon. Until now, running PyTorch workloads on it generally meant Neuron-specific code paths, which raised the switching cost for teams already standardized on Nvidia GPUs and the CUDA ecosystem.

Between the lines

The disclosed test data suggests how much of PyTorch the first release actually covers: roughly 22,000 skipped tests trace to operators with no native Neuron lowering, and five feature families were excluded outright. Meta PyTorch engineering manager Yana van Grunen framed the value as consistent developer experience across hardware platforms, and the one-line port is aimed squarely at that switching cost rather than at raw performance claims.

What's next

Watch the public preview from 20 October and the PyTorch Conference sessions in San Jose on 20-21 October, plus the next release, when AWS says TorchAO and FP8 support move to general availability.

Source: PyTorch Blog (AWS)