This document is relevant for: Trn1, Trn2, Trn3

Training paths (native and legacy XLA)#

Train models on AWS Trainium with PyTorch. This section splits by the PyTorch path you build on: native PyTorch (TorchNeuron), or the legacy PyTorch/XLA path on the older torch-neuronx XLA library. Use TorchNeuron for new work. The XLA path stays documented for workloads already running on it.

Warning

TorchNeuron and the legacy torch-neuronx (XLA) library share the torch-neuronx package name but are different products with different APIs and version lineages. Pick one path and stay on it; do not mix their APIs in a single workload.

Native PyTorch training (TorchNeuron)#

TorchNeuron is the Neuron native PyTorch backend. It runs eager mode, torch.compile with the neuron backend, and the standard distributed APIs: FSDP, DTensor, DDP, and functional collectives. Most PyTorch training code runs on Trainium after you change device placement from cuda to neuron.

See Training best practices for choosing FSDP, tensor parallelism, pipeline parallelism, and meshes on Trainium.

TorchNeuron (native PyTorch on Neuron) overview

Section landing: overview, quickstarts, tutorials, guides, and reference.

TorchNeuron training tutorials tutorial

Worked examples, from a single training loop to distributed training with torch.distributed.

Distributed training how-to how-to

Use FSDP, DTensor, DDP, and functional collectives on Neuron.

Supported Training Libraries#

TorchTitan Neuron overview

Pre-train on Trainium with TorchTitan, including Llama, DeepSeek (FSDP), and GPT-OSS.

Legacy PyTorch/XLA training (torch-neuronx XLA)#

Note

The torch-neuronx library referenced here is the older XLA-based product, not TorchNeuron. It is in maintenance status for customers with existing XLA workloads. For new training, use the native PyTorch path above.

Training (torch-neuronx XLA) overview

Entry point for the legacy XLA-based training docs on Trainium.

NeuronX Distributed Training (NxDT) overview

The XLA-era distributed training library, now in maintenance status and archived.

NeuronX Distributed Core (NxD Core) overview

Distributed-computing primitives the XLA training path builds on.

This document is relevant for: Trn1, Trn2, Trn3