This document is relevant for: Inf2, Trn1, Trn2, Trn3

NeuronX Distributed Setup

Install PyTorch Neuron on Trn1 to create a pytorch environment. It is recommended to work out of a Python virtual environment (such as venv) so as to avoid package installation issues.

You can install the neuronx-distributed package using the following command:

python -m pip install neuronx_distributed --extra-index-url https://pip.repos.neuron.amazonaws.com

This document is relevant for: Inf2, Trn1, Trn2, Trn3