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

update args sanity check for cuda memory balance

pull/306/head
mwiacx 2023-09-22 10:10:58 +08:00
parent 69d542f288
commit 64f0d838e2
1 changed files with 2 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
internlm/initialize/launch.py
View File

@ -292,6 +292,8 @@ def args_sanity_check():
optim_ckpt._add_item("overlap_sync_grad", False) optim_ckpt._add_item("overlap_sync_grad", False)
if "overlap_sync_param" not in optim_ckpt: if "overlap_sync_param" not in optim_ckpt:
optim_ckpt._add_item("overlap_sync_param", False) optim_ckpt._add_item("overlap_sync_param", False)
if "cuda_memory_balance" not in optim_ckpt:
optim_ckpt._add_item("cuda_memory_balance", False)
if gpc.is_rank_for_log(): if gpc.is_rank_for_log():
logger.info( logger.info(
f"overlap_sync_grad:{optim_ckpt.overlap_sync_grad}, overlap_sync_param:{optim_ckpt.overlap_sync_param}" f"overlap_sync_grad:{optim_ckpt.overlap_sync_grad}, overlap_sync_param:{optim_ckpt.overlap_sync_param}"