This document is relevant for: Trn2, Trn3

torch.compile on Neuron

TorchNeuron provides native torch.compile support on AWS Trainium through a custom TorchDynamo backend, selected with torch.compile(model, backend="neuron"). This section explains how compilation works on Neuron and how to use it to optimize your PyTorch models.

How torch.compile works deep dive

The compilation pipeline (TorchDynamo → AOT Autograd → StableHLO → neuronx-cc → NEFF), caching, graph breaks, and fallback behavior.

Optimize a model how-to

Apply backend="neuron" to a working eager model and tune for throughput.

Backend reference reference

Environment variables, the neuron_dynamo_backend helper API, and caching configuration.

This document is relevant for: Trn2, Trn3