Exporting mistral model to neuron fails with ModuleNotFoundError: No module named 'optimum.commands.env'

[thesyntaxinator]

System Info

optimum-neuron            0.0.18
optimum                   1.16.2
platform: amazon linux 2
Python 3.8.16

Traceback (most recent call last):
  File "[venv path]/optimum-cli", line 5, in <module>
    from optimum.commands.optimum_cli import main
  File "[venv path]/lib64/python3.8/site-packages/optimum/commands/__init__.py", line 16, in <module>
    from .env import EnvironmentCommand
ModuleNotFoundError: No module named 'optimum.commands.env'

Who can help?

No response

Information

• [X] The official example scripts
• [ ] My own modified scripts

Tasks

• [ ] An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• [X] My own task or dataset (give details below)

Reproduction (minimal, reproducible, runnable)

Install python, neuron, optimum, optimum neuron versions above on amazon linux 2. Try to export a local fine-tuned mistral model to neuron using any --task (I tried text-classification)

command I used: optimum-cli export neuron --model base_model --task text-classification --sequence_length 4096 --batch_size 1 pytorch_model_neuron/

Expected behavior

Model should export to neuron

[chintanckg]

Please specify the underlying hardware (inferentia2/inferentia1/trainium etc.)

[thesyntaxinator]

I trying to compile on AL2 non inferentia instance to inf1 or inf2

[chintanckg]

I am slightly confused by your reply, what do you mean by "non inferentia instance to inf1 or inf2"?

Basically, if you want to compile a neuron-supported model for inf1 or inf2, the compilation needs to happen on inf1 or inf2 or tranium.

[thesyntaxinator]

I see, will try compiling on inf1 or inf2 and post an update

[thesyntaxinator]

Compiling on inf2 worked, thanks