torchrun --nproc_per_node 2 examples/puzzletron/main.py fails with ModuleNotFoundError

[danielkorzekwa]

Modelopt branch: release/0.44.0

Calling torchrun --nproc_per_node 2 examples/puzzletron/main.py --config examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/llama-3_1-8B_pruneffn_memory.yaml 2>&1 | tee ./log.txt | grep "Puzzletron Progress" fails with

File "/workspace/Model-Optimizer/examples/puzzletron/main.py", line 40, in <module> import modelopt.torch.puzzletron as mtpz ModuleNotFoundError: No module named 'modelopt.torch.puzzletron'

notes:

• calling export PYTHONPATH=$PYTHONPATH:/workspace/Model-Optimizer first, solves this problem

[kevalmorabia97]

can you try pip uninstall -y nvidia-modelopt as well? Nemo containers have 2 modelopt installations (system python and uv's venv) so re-installing in venv causes other system installation to take priority in path causing the import failure