Update Neuron initializations (#7952)

test/neuron/run_tests.sh

@@ -0,0 +1,4 @@
+#!/bin/bash
+set -xue
+
+python3  test/neuron/test_neuron_utils.py

test/neuron/test_neuron_utils.py

@@ -0,0 +1,57 @@
+import os
+import pytest
+import unittest
+from torch_xla._internal.neuron_utils import *
+
+
+class NeuronUtilsTest(unittest.TestCase):
+
+  def test_get_visible_cores_list(self):
+    os.environ["NEURON_RT_VISIBLE_CORES"] = "1"
+    assert (get_visible_cores_list() == [1])
+    os.environ["NEURON_RT_VISIBLE_CORES"] = "1,2,3"
+    assert (get_visible_cores_list() == [1, 2, 3])
+    os.environ["NEURON_RT_VISIBLE_CORES"] = "1-3"
+    assert (get_visible_cores_list() == [1, 2, 3])
+    os.environ["NEURON_RT_VISIBLE_CORES"] = "1-3,5-8"
+    assert (get_visible_cores_list() == [1, 2, 3, 5, 6, 7, 8])
+    os.environ["NEURON_RT_VISIBLE_CORES"] = "1,3,5-8"
+    assert (get_visible_cores_list() == [1, 3, 5, 6, 7, 8])
+    os.environ["NEURON_RT_VISIBLE_CORES"] = "1-3,5-8,3-5"
+    with pytest.raises(ValueError):
+      get_visible_cores_list()
+    os.environ["NEURON_RT_VISIBLE_CORES"] = "1-3,5-8-5"
+    with pytest.raises(ValueError):
+      get_visible_cores_list()
+    os.environ["NEURON_RT_VISIBLE_CORES"] = "a-b,5-8-5"
+    with pytest.raises(Exception):
+      get_visible_cores_list()
+    os.environ["NEURON_RT_VISIBLE_CORES"] = "a"
+    with pytest.raises(Exception):
+      get_visible_cores_list()
+
+  def test_remap_visible_cores(self):
+    os.environ["NEURON_RT_VISIBLE_CORES"] = "1"
+    remap_visible_cores(0, 1)
+    assert (os.environ['NEURON_RT_VISIBLE_CORES'] == "1")
+    os.environ["NEURON_RT_VISIBLE_CORES"] = "1,2,3"
+    remap_visible_cores(1, 3)
+    assert (os.environ['NEURON_RT_VISIBLE_CORES'] == "2")
+    os.environ["NEURON_RT_VISIBLE_CORES"] = "1-3"
+    remap_visible_cores(2, 3)
+    assert (os.environ['NEURON_RT_VISIBLE_CORES'] == "3")
+    os.environ["NEURON_RT_VISIBLE_CORES"] = "1-3,5-8"
+    remap_visible_cores(5, 7)
+    assert (os.environ['NEURON_RT_VISIBLE_CORES'] == "7")
+    os.environ["NEURON_RT_VISIBLE_CORES"] = "1,3,5-8"
+    remap_visible_cores(5, 6)
+    assert (os.environ['NEURON_RT_VISIBLE_CORES'] == "8")
+    with pytest.raises(ValueError):
+      remap_visible_cores(5, 9)
+    with pytest.raises(ValueError):
+      remap_visible_cores(6, 6)
+
+
+if __name__ == "__main__":
+  test = unittest.main()
+  sys.exit(0 if test.result.wasSuccessful() else 1)

test/pjrt/test_runtime.py

@@ -86,7 +86,7 @@ def test_pjrt_default_device(self, env_vars, expect_using_pjrt):
         reload(torch_xla)
         logs_context = contextlib.nullcontext()
         if expect_using_pjrt:
-          self.assertIn(xr.device_type(), ['CPU', 'CUDA', 'TPU'])
+          self.assertIn(xr.device_type(), ['CPU', 'CUDA', 'TPU', 'NEURON'])
         else:
           self.assertIsNone(xr.device_type())
 

torch_xla/__init__.py

@@ -116,11 +116,31 @@ def _summarize_fn_tracker():
 
 def _aws_ec2_inf_trn_init():
   try:
-    from torch_neuronx import xla
+    from libneuronxla.libneuronpjrt_path import libneuronpjrt_path
   except ImportError:
-    return
+    pass
   else:
-    xla.init()
+    # Need to set NEURON_LIBRARY_PATH here for proper Neuron Cache behavior
+    os.environ.setdefault('NEURON_LIBRARY_PATH', libneuronpjrt_path())
+    # Enable addition features and overrides
+    try:
+      from torch_neuronx import xla
+    except ImportError:
+      pass
+    else:
+      xla.init()
+
+    # Basic initializations if torch-neuronx is not available
+    from ._internal import neuron
+    if os.path.basename(sys.argv[0]) != 'neuron_parallel_compile':
+      import libneuronxla
+      libneuronxla.configure_environment()
+      neuron.set_envvar_defaults()
+      neuron.configure_pjrt_environment()
+    # Found libneuronxla
+    return True
+  # Did not find libneuronxla
+  return False
 
 
 def _setup_tpu_vm_library_path() -> bool:
@@ -179,7 +199,7 @@ def _check_deprecated_env_var():
 _found_libtpu = _setup_tpu_vm_library_path()
 
 # Setup Neuron library for AWS EC2 inf/trn instances.
-_aws_ec2_inf_trn_init()
+_found_libneuronxla = _aws_ec2_inf_trn_init()
 
 
 def _prepare_to_exit():

torch_xla/_internal/neuron.py

@@ -3,33 +3,115 @@
 
 from torch_xla.experimental import plugins
 
+import sys
+import torch.distributed as dist
+
+from .neuron_utils import get_visible_cores_list, remap_visible_cores
+
+logging.basicConfig()
+logger = logging.getLogger(__name__)
+
+
+# Set root communication address/port
+def set_rt_root_comm_id():
+  if os.environ.get('NEURON_RT_ROOT_COMM_ID', None) is None:
+    if 'MASTER_ADDR' not in os.environ:
+      logging.warning(
+          "MASTER_ADDR environment variable is not set, defaulting to localhost"
+      )
+    root_port = 62182
+    root_addr = os.environ.get('MASTER_ADDR', 'localhost')
+    is_ipv6 = len(root_addr.split(":")) >= 3
+    if is_ipv6:
+      modified = False
+      if not root_addr.startswith("["):
+        root_addr = "[" + root_addr
+        modified = True
+      if not root_addr.endswith("]"):
+        root_addr = root_addr + "]"
+        modified = True
+      if modified:
+        logger.warning(
+            "IPv6 address detected for MASTER_ADDR and missing brackets added: {}"
+            .format(root_addr))
+    os.environ['NEURON_RT_ROOT_COMM_ID'] = '{}:{}'.format(root_addr, root_port)
+
+
+def set_envvar_defaults():
+  os.environ.setdefault('ALLREDUCE_GRADIENTS_BUCKET_SIZE_MB', '50')
+
+
+def configure_pjrt_environment():
+  """
+  Setting all necessary PJRT default environment variables.
+  """
+  from torch.distributed import is_torchelastic_launched
+
+  # Set root communication address/port
+  set_rt_root_comm_id()
+
+  # Set env variables if we don't use GSPMD, using PJRT, and using torchrun
+  if os.environ.get('XLA_USE_SPMD', '0') != '1' \
+      and is_torchelastic_launched():
+    # Env variables that only need to be set once
+    # NEURON_PJRT_PROCESSES_NUM_DEVICES is a list of core counts and is too long for very large cluster,
+    # so use NEURON_PJRT_WORLD_SIZE to pass world size and use core count of 1 per process in PJRT client.
+    if 'NEURON_PJRT_PROCESSES_NUM_DEVICES' not in os.environ and 'NEURON_PJRT_WORLD_SIZE' not in os.environ:
+      if 'WORLD_SIZE' not in os.environ:
+        logger.warning(
+            'WORLD_SIZE environment variable not set, defaulting to 1.')
+      os.environ["NEURON_PJRT_WORLD_SIZE"] = os.environ.get("WORLD_SIZE", "1")
+      if 'LOCAL_WORLD_SIZE' not in os.environ:
+        logger.warning(
+            'LOCAL_WORLD_SIZE environment variable not set, defaulting to 1.')
+      os.environ['PJRT_LOCAL_PROCESS_COUNT'] = os.environ.get(
+          'LOCAL_WORLD_SIZE', '1')
+
+    # Env variables that need to be set once per process
+    if not os.environ.get('NEURON_RT_VISIBLE_CORES', None):
+      os.environ['NEURON_RT_VISIBLE_CORES'] = os.environ.get('LOCAL_RANK', '0')
+    else:
+      local_rank = int(os.environ.get('LOCAL_RANK', '0'))
+      local_world_size = int(os.environ.get('LOCAL_WORLD_SIZE', '1'))
+      remap_visible_cores(local_rank, local_world_size)
+
+    if 'RANK' not in os.environ:
+      logger.warning('RANK environment variable is not set, defaulting to 0.')
+    os.environ['NEURON_PJRT_PROCESS_INDEX'] = os.environ.get('RANK', '0')
+    if 'LOCAL_RANK' not in os.environ:
+      logger.warning(
+          'LOCAL RANK environment variable is not set, defaulting to 0.')
+    os.environ['PJRT_LOCAL_PROCESS_RANK'] = os.environ.get('LOCAL_RANK', '0')
+
 
 def num_local_processes() -> int:
-  if 'MASTER_ADDR' not in os.environ:
-    logging.warning("MASTER_ADDR not setting, defaulting to localhost")
-  os.environ['NEURON_RT_ROOT_COMM_ID'] = '{}:{}'.format(
-      os.environ.get('MASTER_ADDR', 'localhost'), '62182')
-  if "NEURONCORE_NUM_DEVICES" not in os.environ:
-    logging.warning("NEURONCORE_NUM_DEVICES not set, defaulting to 1")
+  set_rt_root_comm_id()
   num_processes = int(os.environ.get("NEURONCORE_NUM_DEVICES", "1"))
   os.environ['NEURON_PJRT_PROCESSES_NUM_DEVICES'] = ','.join(
       ['1' for _ in range(num_processes)])
-
   return num_processes
 
 
+# When torchrun is used, setting these environments causes the
+# second instance in 2-node cluster to think it is node 0 instead of node 1.
+# Need to skip these settings and let configure_pjrt_environment to
+# set the distributed PJRT environment variables.
+# If NEURONCORE_NUM_DEVICES is used, then go ahead and set the environments.
 def initialize_env(local_rank, local_world_size):
-  os.environ["NEURON_PJRT_PROCESS_INDEX"] = str(local_rank)
-  assert (
-      local_rank < local_world_size
-  ), "ERROR in initialize_env: PJRT_LOCAL_PROCESS_RANK is not less than PJRT_LOCAL_PROCESS_COUNT"
-  os.environ["NEURON_RT_VISIBLE_CORES"] = str(local_rank)
+  from torch.distributed import is_torchelastic_launched
+  if not is_torchelastic_launched():
+    os.environ["NEURON_PJRT_PROCESS_INDEX"] = str(local_rank)
+    if not os.environ.get('NEURON_RT_VISIBLE_CORES', None):
+      os.environ["NEURON_RT_VISIBLE_CORES"] = str(local_rank)
+    else:
+      remap_visible_cores(local_rank, local_world_size)
 
 
 class NeuronPlugin(plugins.DevicePlugin):
 
   def library_path(self):
-    return os.environ.get("NEURON_LIBRARY_PATH", "libneuronpjrt.so")
+    from libneuronxla.libneuronpjrt_path import libneuronpjrt_path
+    return os.environ.get("NEURON_LIBRARY_PATH", libneuronpjrt_path())
 
   def configure_multiprocess(self, local_rank, local_world_size):
     initialize_env(local_rank, local_world_size)

torch_xla/_internal/neuron_utils.py

@@ -0,0 +1,66 @@
+import os
+import logging
+logging.basicConfig()
+logger = logging.getLogger(__name__)
+
+
+def convert_range(range_spec):
+  try:
+    lowerupper = list(map(int, range_spec.split("-")))
+  except Exception as e:
+    print(f"ERROR: Malformed range specs in NEURON_RT_VISIBLE_CORES;" +
+          f"expecting <int> or <lower int>-<upper int> (got {range_spec})")
+    raise e
+  if len(lowerupper) > 2:
+    raise ValueError(
+        f"ERROR: Range specs in NEURON_RT_VISIBLE_CORES should be of " +
+        f"the form <int> or <lower int>-<upper int> (got {range_spec})")
+  if len(lowerupper) == 2:
+    if lowerupper[0] > lowerupper[1]:
+      raise ValueError(
+          f"ERROR: Range specs in NEURON_RT_VISIBLE_CORES should " +
+          f"be of the form <int> or <lower int>-<upper int> (got {range_spec})")
+    lowerupper = range(lowerupper[0], lowerupper[1] + 1)
+  return lowerupper
+
+
+def get_visible_cores_list():
+  import os
+
+  range_list = os.environ.get("NEURON_RT_VISIBLE_CORES")
+  cores_list = None
+  if range_list:
+    range_list = range_list.split(",")
+    cores_list = []
+    for i in range_list:
+      new = convert_range(i)
+      if (set(cores_list) & set(new)) != set():
+        raise ValueError(
+            "ERROR: Please ensure the ranges in NEURON_RT_VISIBLE_CORES are mutually exclusive."
+        )
+      cores_list.extend(new)
+  return cores_list
+
+
+def remap_visible_cores(local_rank, local_world_size):
+  cores_list = get_visible_cores_list()
+  count = len(cores_list)
+  assert (local_world_size > 0), "Local world size should be non-zero"
+  if count <= 1 and local_world_size == 1:
+    # Allow user to pass NEURON_RT_VISIBLE_CORES for sinlge-core workload
+    pass
+  elif local_world_size != count:
+    raise ValueError(
+        f"LOCAL_WORLD_SIZE (torchrun) or PJRT_LOCAL_PROCESS_COUNT (xmp.spawn) value of {local_world_size} "
+        +
+        f"is not equal to count {count} from NEURON_RT_VISIBLE_CORES {cores_list}"
+    )
+  elif local_rank >= count:
+    raise ValueError(
+        f"LOCAL_RANK (torchrun) or PJRT_LOCAL_PROCESS_RANK (xmp.spawn) value of {local_rank} is higher than "
+        + f"count {count} from NEURON_RT_VISIBLE_CORES {cores_list}")
+  else:
+    remapped_core = cores_list[local_rank]
+    logger.warning(f"Remapping NEURON_RT_VISIBLE_CORES {cores_list} to " +
+                   f"NEURON_RT_VISIBLE_CORES[LOCAL_RANK]={remapped_core}")
+    os.environ['NEURON_RT_VISIBLE_CORES'] = str(remapped_core)

torch_xla/runtime.py

@@ -72,6 +72,9 @@ def _maybe_select_default_device():
         + num_devices_str)
     os.environ[xenv.PJRT_DEVICE] = 'CUDA'
     os.environ[xenv.GPU_NUM_DEVICES] = num_devices_str
+  elif torch_xla._found_libneuronxla:
+    logging.warning('Found libneuronpjrt.so. Setting PJRT_DEVICE=NEURON.')
+    os.environ[xenv.PJRT_DEVICE] = 'NEURON'
   else:
     logging.warning('Defaulting to PJRT_DEVICE=CPU')
     os.environ[xenv.PJRT_DEVICE] = 'CPU'