Function _get_expert_parallel_ranks

deepspeed/utils/groups.py:385–451

Generate expert parallel and expert data parallel group ranks list. Example - E + M + D parallel world_size = 16 model_degree = 2 expert_degree = 4 # number of experts in same group mp_group = [0, 1], [2,3], [4,5] ... data_parallel_group =[0,2,4,6,8,1

(world_size,
                               tensor_parallel_size_,
                               expert_parallel_size_,
                               pipeline_parallel_size_=1,
                               use_data_before_expert_parallel_=False)

def _get_expert_parallel_ranks(world_size,
                               tensor_parallel_size_,
                               expert_parallel_size_,
                               pipeline_parallel_size_=1,
                               use_data_before_expert_parallel_=False):
    """Generate expert parallel and expert data parallel group ranks list.

        Example - E + M + D parallel
        world_size = 16
        model_degree = 2
        expert_degree = 4 # number of experts in same group
        mp_group = [0, 1], [2,3], [4,5] ...
        data_parallel_group =[0,2,4,6,8,10, 12,14],                 [1,3,5,7,9,11,13,15]
        expert_parallel_group = [0,2,4,6], [8,10,12,14]             [1,3,5,7], [9,11,13,15]
        expert_data_parallel_group = [0,8],[2,10],[4,12],[6,14],    [1,9],[3,11],[5,13],[7,15]

    Args:
        world_size (int): Distributed world size.
        tensor_parallel_size_ (int): Tensor parallel group size.
        expert_parallel_size_ (int): Expert parallel group size.
        pipeline_parallel_size_ (int): Pipeline parallel group size
        use_data_before_expert_parallel_ (bool): Use the D + E instead of E + D topology
    Returns:
        Expert parallel group ranks and Expert data parallel group ranks list.
    """
    _ensure_divisibility(world_size, tensor_parallel_size_ * pipeline_parallel_size_)
    dp_world_size = world_size // (tensor_parallel_size_ * pipeline_parallel_size_)
    _ensure_divisibility(dp_world_size, expert_parallel_size_)

    # Generate data parallel groups
    data_parallel_groups = []
    dp_group_size = tensor_parallel_size_
    pp_stride = world_size // pipeline_parallel_size_

    if use_data_before_expert_parallel_:
        dp_stride = world_size // expert_parallel_size_ // tensor_parallel_size_ // pipeline_parallel_size_
        for pp_stage_start in range(0, world_size, pp_stride):
            pp_stage_next = pp_stage_start + pp_stride
            for i in range(dp_group_size):
                data_parallel_groups.append(list())
                for ds in range(dp_stride):
                    # [0, 4, 8, 12, 16, 20, 24, 28, 2, 6, 10, 14, 18, 22, 26, 30]
                    # [1, 5, 9, 13, 17, 21, 25, 29, 3, 7, 11, 15, 19, 23, 27, 31]
                    data_parallel_groups[-1].extend(
                        list(
                            range(pp_stage_start + i + ds * tensor_parallel_size_, pp_stage_next,
                                  dp_stride * tensor_parallel_size_)))
    else:
        for pp_stage_start in range(0, world_size, pp_stride):
            pp_stage_next = pp_stage_start + pp_stride
            for i in range(dp_group_size):
                data_parallel_groups.append(list(range(pp_stage_start + i, pp_stage_next, dp_group_size)))

    expert_parallel_groups = []
    expert_data_parallel_groups = []
    for dp_ranks in data_parallel_groups:
        # partition of expert parallel groups, e.g. [0,2,4,6], [8,10,12,14]
        part_ep_groups = []