From 3e6e082fa5118e113ab18bc548151940f23911c0 Mon Sep 17 00:00:00 2001 From: li126com Date: Fri, 22 Sep 2023 18:15:41 +0800 Subject: [PATCH] feat: add numa --- internlm/initialize/launch.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/internlm/initialize/launch.py b/internlm/initialize/launch.py index ee0cbad..2a1ccbc 100644 --- a/internlm/initialize/launch.py +++ b/internlm/initialize/launch.py @@ -478,15 +478,18 @@ def get_config_value(config, key, defalut): def try_bind_numa(launcher): + # check pacakge try: import numa from numa import memory, schedule from pynvml.smi import nvidia_smi except (AttributeError, ImportError): - logger.info( - "Try bind numa failed! Package import error, if numa is not installed, " - "please implement: pip install --upgrade py-libnuma" - ) + global_rank = int(os.environ["SLURM_PROCID"]) if launcher == "slurm" 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" + ) return # get numa node number