PyTorch NeuronX Application Notes
This section contains application notes specific to PyTorch NeuronX (torch-neuronx
) for Trn1
and Inf2
instances. These guides cover PyTorch version migrations, advanced features, optimization techniques, and best practices for training and inference on AWS Trainium and Inferentia2.
PyTorch Version Support
New features and migration guide for PyTorch 2.8 on Neuron
Features and improvements introduced with PyTorch 2.7 support
Migration guide and new capabilities for PyTorch 2.6
General guide to PyTorch 2.x series support and features
Advanced Features
Advanced graph partitioning strategies for distributed training and inference
Scale inference workloads using torch.neuronx.DataParallel
for multi-core execution
Migrate from deprecated XLA environment variables to PyTorch mixed-precision options