Neuron model NEFFs are dependent on the python path

[dacorvo]

The same bug that was present in AWS Neuron SDK 2.19 and fixed in 2.19.1 (#91) is back in AWS Neuron SDK 2.20.

With AWS Neuron SDK 2.19, when exporting a model and saving the compiled artifacts, it is impossible to reload them afterwards if the python path is different.

This basically makes shared serialization and caching impossible, since you cannot control the deployment environment (ec2 with DLAMI, sagemaker or ad-hoc end-user endpoints will all have different environments).

Steps to reproduce

1. download test_tnx_llama_export.py

2. export the model in a venv

$ python3 -m venv foo_venv
$ source foo_venv/bin/activate
$ export PIP_EXTRA_INDEX_URL=https://pip.repos.neuron.amazonaws.com
$ python - m pip install -U neuronx-cc torch_neuronx==2.* transformers-neuronx
$ python test_tnx_llama_export.py export meta-llama/Llama-3.1-8B-Instruct --save_dir ./llama-foo

3. check the generated artifacts and verify the neuron model can be reloaded (no compilation should happen)

$ python test_tnx_llama_export.py run meta-llama/Llama-3.1-8B-Instruct --save_dir ./llama-foo

4. deactivate the venv and try to reload the model in another venv

$ deactivate
$ python3 -m venv bar_venv
$ source bar_venv/bin/activate
$ export PIP_EXTRA_INDEX_URL=https://pip.repos.neuron.amazonaws.com
$ python -m pip install -U neuronx-cc torch_neuronx==2.* transformers-neuronx
$ python test_tnx_llama_export.py run meta-llama/Llama-3.1-8B-Instruct --save_dir ./llama-foo

You should get the following exception:

FileNotFoundError: Could not find a matching NEFF for your HLO in this directory. Ensure that the model you are trying to load is the same type and has the same parameters as the one you saved or call "save" on this model to reserialize it.

5. export the model from the new venv

$ python test_tnx_llama_export.py export meta-llama/Llama-3.1-8B-Instruct --save_dir ./llama-bar

Now if you compare the NEFF files in the two save dir you will see that one of them is different.

[pagezyhf]

@jeffhataws as you helped working on this last time (2.19 to 2.19.1)?

[aws-patlange]

Thank you for reporting this bug. We have reproduced and identified the issue and are working on a fix.