aws-neuron/aws-neuron-sdk

Bug in `configure_pjrt_environment`

michaelbenayoun opened this issue · 1 comments

There is a bug in configure_pjrt_environment (torch_neuronx/xla.py).

        if not os.environ.get('NEURON_RT_VISIBLE_CORES', None):
            os.environ['NEURON_RT_VISIBLE_CORES'] = os.environ.get('LOCAL_RANK', 0)

Should be:

        if not os.environ.get('NEURON_RT_VISIBLE_CORES', None):
            os.environ['NEURON_RT_VISIBLE_CORES'] = os.environ.get('LOCAL_RANK', '0')

Without that, when no environment variable is set, the system crashes.

Hi @michaelbenayoun,

Thank you for filing this issue, we had not caught this previously.

We will apply your fix for the issue and will include it in an upcoming release.