libneuronxla version is not pinned in torch-neuronx, causing unexpected graph changes and error(s) like "Unknown flags in XLA_FLAGS"

[jeffhataws]

Currently, recent versions of torch-neuronx has a loose dependency on libneuronxla:

libneuronxla>2.0.965,<3.0

When trying to install torch-neuronx from older release (say SDK 2.19.1), it will bring in newer version of libneuronxla (say from SDK 2.20). For example, if installing torch-neuronx==2.1.2.2.2.0 (SDK 2.19.1) it will install libneuronxla version 2.0.4115.0 instead of 2.0.2335 (inconsistent versions):

libneuronxla                            2.0.4115.0
torch-neuronx                           2.1.2.2.2.0
torch-xla                               2.1.3

This can cause graph cache misses when using optimum-neuron's cache for example. Furthermore, the release 2.20 libneuronxla also has support for newer openxla and thus you may get an error "Unknown flags in XLA_FLAGS" like below:

F external/xla/xla/parse_flags_from_env.cc:224] Unknown flags in XLA_FLAGS: --xla_gpu_simplify_all_fp_conversions=false  --xla_gpu_force_compilation_parallelism=8 "

The workaround is to pin to the released version based on the 2.19.1 released artifacts.

python -m pip install neuronx-cc==2.* torch-neuronx==2.1.2.2.2.0 libneuronxla==2.0.2335 torchvision