Installation

Our recommendation is to use the Hugging Face Neuron Deep Learning AMI (DLAMI). The DLAMI comes with all required libraries pre-packaged for you, including the Optimum Neuron, Neuron Drivers, Transformers, Datasets, and Accelerate.

But it can also be installed using pip as described below.

Before installing optimum-neuron make sure that you have installed the Neuron driver and tools, check out more detailed guide here.

Adding pip packages URL

Pointing to the AWS Neuron repository:

python -m pip config set global.extra-index-url https://pip.repos.neuron.amazonaws.com

Installing optimum-neuron for AWS Trainium ( trn1 ) or AWS inferentia2 ( inf2 )

python -m pip install --upgrade-strategy eager optimum[neuronx]

Installing optimum-neuron for AWS inferentia ( inf1 )

python -m pip install --upgrade-strategy eager optimum[neuron]