No module named 'optimum.neuron'; 'optimum' is not a package

I am trying to use Optimum for Neuron on an AWS Inf2 instance.
I installed optimum:

python3 -m pip install git+[neuron]

Everything seems to be installed correctly but when I try to import optimum.neuron I get “ModuleNotFoundError”:

from optimum.neuron import NeuronModelForCausalLM

ModuleNotFoundError: No module named ‘optimum.neuron’; ‘optimum’ is not a package

Anyone know what is the correct way to import the package?

Hi @ria14313, I cannot reproduce your issue. But maybe you can try following the installation instructions provided in the documentation: Installation

Hi @ria14313,

If you are using inf2, the extra for installing optimum-neuron package should be neuronx instead of neuron (for inf1).

Normally the neuronx extra will help you to install the neuronx-cc compiler, if not you can follow the guide of neuronx installation from AWS.

Make sure that you have the following requirements:

  • neuronx driver and runtime
  • neuronx-cc==2.* torch-neuronx

Then you can install optimum with the installation guide.

1 Like