init

trainer_misc/sp_utils.py

@@ -0,0 +1,98 @@
+import os
+import torch
+import torch.distributed as dist
+from .utils import is_dist_avail_and_initialized, get_rank
+
+
+SEQ_PARALLEL_GROUP = None
+SEQ_PARALLEL_SIZE = None
+SEQ_PARALLEL_PROC_NUM = None    # using how many process for sequence parallel
+
+SYNC_INPUT_GROUP = None
+SYNC_INPUT_SIZE = None
+
+def is_sequence_parallel_initialized():
+    if SEQ_PARALLEL_GROUP is None:
+        return False
+    else:
+        return True
+
+
+def init_sequence_parallel_group(args):
+    global SEQ_PARALLEL_GROUP
+    global SEQ_PARALLEL_SIZE
+    global SEQ_PARALLEL_PROC_NUM
+
+    assert SEQ_PARALLEL_GROUP is None, "sequence parallel group is already initialized"
+    assert is_dist_avail_and_initialized(), "The pytorch distributed should be initialized"
+    SEQ_PARALLEL_SIZE = args.sp_group_size
+    
+    print(f"Setting the Sequence Parallel Size {SEQ_PARALLEL_SIZE}")
+
+    rank = torch.distributed.get_rank()
+    world_size = torch.distributed.get_world_size()
+
+    if args.sp_proc_num == -1:
+        SEQ_PARALLEL_PROC_NUM = world_size
+    else:
+        SEQ_PARALLEL_PROC_NUM = args.sp_proc_num
+
+    assert SEQ_PARALLEL_PROC_NUM % SEQ_PARALLEL_SIZE == 0, "The process needs to be evenly divided"
+
+    for i in range(0, SEQ_PARALLEL_PROC_NUM, SEQ_PARALLEL_SIZE):
+        ranks = list(range(i, i + SEQ_PARALLEL_SIZE))
+        group = torch.distributed.new_group(ranks)
+        if rank in ranks:
+            SEQ_PARALLEL_GROUP = group
+            break
+
+
+def init_sync_input_group(args):
+    global SYNC_INPUT_GROUP
+    global SYNC_INPUT_SIZE
+
+    assert SYNC_INPUT_GROUP is None, "parallel group is already initialized"
+    assert is_dist_avail_and_initialized(), "The pytorch distributed should be initialized"
+    SYNC_INPUT_SIZE = args.max_frames
+
+    rank = torch.distributed.get_rank()
+    world_size = torch.distributed.get_world_size()
+
+    for i in range(0, world_size, SYNC_INPUT_SIZE):
+        ranks = list(range(i, i + SYNC_INPUT_SIZE))
+        group = torch.distributed.new_group(ranks)
+        if rank in ranks:
+            SYNC_INPUT_GROUP = group
+            break
+
+
+def get_sequence_parallel_group():
+    assert SEQ_PARALLEL_GROUP is not None, "sequence parallel group is not initialized"
+    return SEQ_PARALLEL_GROUP
+
+
+def get_sync_input_group():
+    return SYNC_INPUT_GROUP
+
+
+def get_sequence_parallel_world_size():
+    assert SEQ_PARALLEL_SIZE is not None, "sequence parallel size is not initialized"
+    return SEQ_PARALLEL_SIZE
+
+
+def get_sequence_parallel_rank():
+    assert SEQ_PARALLEL_SIZE is not None, "sequence parallel size is not initialized"
+    rank = get_rank()
+    cp_rank = rank % SEQ_PARALLEL_SIZE
+    return cp_rank
+
+
+def get_sequence_parallel_group_rank():
+    assert SEQ_PARALLEL_SIZE is not None, "sequence parallel size is not initialized"
+    rank = get_rank()
+    cp_group_rank = rank // SEQ_PARALLEL_SIZE
+    return cp_group_rank
+
+
+def get_sequence_parallel_proc_num():
+    return SEQ_PARALLEL_PROC_NUM