This document is relevant for: Trn2, Trn3

TorchNeuron guides

Task-oriented guides for common TorchNeuron workflows on AWS Trainium. Each guide walks you through a specific task with step-by-step instructions.

Work with NKI kernels from torch-neuronx how-to

Integrate NKI kernels into torch-neuronx workflows with @nki_op, autograd, and torch.compile.

Use torch.distributed on Neuron how-to

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

Profile memory usage in torch-neuronx how-to

Use torch_neuronx.memory_* APIs and Neuron Explorer to measure and fix memory issues.

torch.compile on Neuron how-to

JIT-compile your model with backend="neuron". Covers the compilation pipeline, optimization, compiling without hardware, and the backend reference.

This document is relevant for: Trn2, Trn3