feat:add numa

internlm/initialize/launch.py

@@ -477,19 +477,26 @@ def get_config_value(config, key, defalut):
     return value
 
 
-def try_bind_numa(launcher):
+# check pacakge
-    # check pacakge
+try:
-    try:
     import numa
     from numa import memory, schedule
     from pynvml.smi import nvidia_smi
-    except (AttributeError, ImportError):
+except (AttributeError, ImportError):
-        global_rank = int(os.environ["SLURM_PROCID"]) if launcher == "slurm" else int(os.environ["RANK"])
+    get_numa = False
+    global_rank = int(os.environ["SLURM_PROCID"]) if "SLURM_PROCID" in os.environ else int(os.environ["RANK"])
     if global_rank == 0:
         logger.info(
             "Try bind numa failed! Package import error, if numa is not installed, "
             "please implement: pip install --upgrade py-libnuma"
         )
+else:
+    get_numa = True
+
+
+def try_bind_numa(launcher):
+    # Early return if numa module not available
+    if not get_numa:
         return
 
     # get numa node number