Function reinit_dp_communicator

training/comm/comm_utils.py:159–205

(args)

        
        
def reinit_dp_communicator(args):
    
    print('###### reinit start #######')
    
    default_init(args)
    assert args.world_size == args.data_group_size * args.pipeline_group_size
    if args.world_size == args.data_group_size * args.pipeline_group_size:
        #    We do the following hard code alignment of communication groups:
        #    Suppose there are 8 instances (world_size), and 4 data parallel groups (data_group_size is 2),
        #    Then there would be 2 pipeline parallel groups (pipeline_group_size is 4), then the groups will look like:
        #    pipeline parallel: <group 0: [0,1,2,3]>, <group 1: [4,5,6,7]>
        #    data parallel: <group 0: [0,4]>, <group 1: [1,5]>, <group 2: [2,6]>, <group 3: [3,7]>
        # assert args.world_size == args.data_group_size * args.pipeline_group_size
        global _DATA_PARALLEL_COMM
        global _PIPELINE_PARALLEL_COMM
        global _DATA_PARALLEL_RANK
        global _PIPELINE_PARALLEL_RANK
        global _DATA_PARALLEL_WORLD_SIZE
        global _PIPELINE_PARALLEL_WORLD_SIZE
        
        if args.data_group_size != 1:
            
            dp_backend = getattr(args, 'dp_backend', 'gloo')
            if dp_backend == 'nccl':
            
                raise Exception('NCCL cannot reinit.')
            
            elif dp_backend == 'gloo':
                
                for i in range(args.pipeline_group_size):
                    ranks = [rank for rank in range(i, args.world_size, args.pipeline_group_size)]
                    print(args.rank, ranks)
                    data_group = torch.distributed.new_group(ranks, backend='gloo')
                    if args.rank in ranks:
                        def to_global_rank(dp_rank):
                            rank = _PIPELINE_PARALLEL_RANK + dp_rank * args.pipeline_group_size
                            # print(f"{dp_rank} --> {rank}")
                            return rank
                        _DATA_PARALLEL_COMM = TorchCommunicator(
                            data_group, to_global_rank=to_global_rank, 
                            dp_rank=_DATA_PARALLEL_RANK,
                            comm_group_size=args.data_group_size,)
            
            else:
                assert False
            
        print('######## dp comm reinit done!! ########')