This document is relevant for: Trn2, Trn3

TorchNeuron: Native PyTorch on AWS Neuron#

Beta Release

TorchNeuron is in public beta for the Neuron 2.33 release. It is supported on Trainium 2 (trn2) and Trainium 3 (trn3) instances. We welcome your feedback.

TorchNeuron is an open-source native PyTorch backend for AWS Trainium. It integrates through PyTorch’s standard PrivateUse1 device backend mechanism, providing eager mode execution, torch.compile for JIT compilation, and standard distributed APIs including FSDP, DTensor, DDP, and Tensor Parallelism.

TorchNeuron uses standard PyTorch APIs. To run existing PyTorch code on Trainium, change the device placement from cuda to neuron.

What changes are needed?#

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 torch.autocast(device_type="neuron")

  • Distributed training: torch.distributed works natively through the neuron backend and supports the PyTorch parallelism APIs: DDP, FSDP2, tensor and sequence parallelism with DTensor, pipeline parallelism, and context parallelism. Configure sharding based on NeuronCores rather than GPUs

For torch.compile (on top of the eager mode changes):

  • Specify backend="neuron" (e.g., @torch.compile(backend="neuron"))

  • Remove CUDA-specific parameters like mode="max-autotune-no-cudagraphs"

No import torch_neuronx is needed: PyTorch autoloads TorchNeuron on import torch through the torch.backends entry point that the torch-neuronx package registers. This includes @nki.jit kernels, which run without the import. Import torch_neuronx only to call its own APIs, such as nki_op or wrap_nki (see How to work with NKI kernels from torch-neuronx).

About TorchNeuron overview

Understand TorchNeuron’s architecture, execution modes (eager and torch.compile), and how it integrates with the PyTorch ecosystem.

Get started quickstart

Quickstart guides to get your first model running on Trainium in eager mode and with torch.compile.

Tutorials tutorial

End-to-end walk-throughs for simple and distributed training, Hugging Face fine-tuning, and Hugging Face inference with torch.compile.

Guides how-to

How-to guides for torch.distributed on Neuron, working with NKI kernels from torch-neuronx, profiling memory usage, and torch.compile on Neuron.

Native PyTorch inference overview

Run inference with native PyTorch on Trainium, serve native PyTorch models with vLLM Neuron, and see which package each path requires.

TorchTitan Neuron tutorial

Run TorchTitan pre-training on Trainium with TorchNeuron, with tutorials for Llama, DeepSeek (SimpleFSDP), and GPT-OSS.

Reference reference

API compatibility, environment variables, best practices, memory management guidelines, and links to resources.

Troubleshooting troubleshooting

Diagnose and resolve common issues with installation, eager mode, torch.compile, and distributed training.

This document is relevant for: Trn2, Trn3