This document is relevant for: Inf2, Trn1, Trn1n

NeuronX Distributed Setup

Install PyTorch Neuron on Trn1 to create a pytorch environment. It is recommended to work out of python virtual env 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

Make sure the transformers version is set to 4.26.0

This document is relevant for: Inf2, Trn1, Trn1n