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

no overlap for save ckpt

pull/540/head
lijiaxing 2023-12-19 17:49:49 +08:00
parent 9bf24d9768
commit d3ca22cf3d
1 changed files with 1 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
internlm/utils/gputest.py
View File

@ -36,8 +36,7 @@ def empty_cache_and_diag(batch_count, interval=50):
if interval <= 0:
interval = 50
if not gpc.config.hybrid_zero_optimizer.overlap_sync_param:
cuda_memory_analyze(batch_count, batch_count % int(interval) == 0 or batch_count <= 5)
cuda_memory_analyze(batch_count, batch_count % int(interval) == 0 or batch_count <= 5)
if batch_count % int(interval) == 0:
# there is no need to do diag on the first batch
@ -302,7 +301,6 @@ def warmup_process_group():
def cuda_memory_analyze(step=0, print_mm_suage=False):
global n_caching_allocator_flushes
torch.cuda.synchronize()
g_rank = gpc.get_global_rank()
tp_rank = gpc.get_local_rank(ParallelMode.TENSOR)