github
/
InternLM
mirror of https://github.com/InternLM/InternLM
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

feat: add numa

pull/320/head
li126com 2023-09-22 18:15:41 +08:00
parent cde4c57ce9
commit 3e6e082fa5
1 changed files with 7 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
internlm/initialize/launch.py
View File

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