This document is relevant for: Trn2, Trn3

TorchNeuron resources#

Links to the source code, examples, profiling tools, upstream PyTorch documentation, and support channels for TorchNeuron (torch-neuronx).

Source code#

aws-neuron/torch-neuronx contains the TorchNeuron native PyTorch backend: the PrivateUse1 device backend, ATen operator implementations, and the TorchDynamo backend for torch.compile.

Examples#

The examples/ directory contains ready-to-run scripts, each with its own README. Start with the torch.compile quickstart, the GPT-2 training loop in eager mode, fused attention with NKI kernels, or tensor parallelism with FSDP.

Profile with Neuron Explorer#

Neuron Explorer: what Neuron Explorer shows about a TorchNeuron workload.
Capture profiles with Neuron Explorer: capture a profile from a PyTorch workload, including torch.compile versus eager mode.

Libraries built on TorchNeuron#

vLLM Neuron and vLLM Omni Neuron for inference serving.
TorchTitan for large-scale pre-training.
Hugging Face transformers, diffusers, PEFT, and accelerate run natively on TorchNeuron; no Neuron-specific compatibility layer is needed.

PyTorch community resources#

torch.compile overview: torch.compile and TorchDynamo.
Custom backends for torch.compile: how a backend such as backend="neuron" plugs into torch.compile.
PrivateUse1 backend tutorial: the out-of-tree device mechanism TorchNeuron uses.

Get help#

Troubleshooting: TorchNeuron: common issues and what to collect before reporting a problem.
aws-neuron/aws-neuron-sdk issues: report bugs and request features.
AWS re:Post, aws-neuron tag: ask questions and search community answers.

Migrating from XLA-based torch-neuronx#

Earlier torch-neuronx releases were built on PyTorch/XLA and are a separate product lineage from TorchNeuron. The XLA-based documentation is in the archive.