Function _create_expert_data_and_model_parallel

deepspeed/utils/groups.py:454–508

Create expert and data parallel groups based on MPU (model parallel) group. Note: Caller of this function is responsible to check if the groups already exist. Example - E + M + D parallel world_size = 16 model_degree = 2 expert_degree = 4 # number o

(expert_parallel_size_, mpu, use_data_before_expert_parallel_=False)

def _create_expert_data_and_model_parallel(expert_parallel_size_, mpu, use_data_before_expert_parallel_=False):
    """
        Create expert and data parallel groups based on MPU (model parallel) group.

        Note: Caller of this function is responsible to check if the groups already exist.

        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]
    """
    assert dist.is_initialized(), "dist is not initialized"
    tensor_parallel_size_ = bwc_tensor_model_parallel_world_size(mpu)

    global expert_tensor_parallel_world_size
    expert_tensor_parallel_world_size = tensor_parallel_size_

    world_size = dist.get_world_size()
    rank = dist.get_rank()
    dp_world_size = _get_data_parallel_world_size()
    pp_world_size = 1 if mpu is None else bwc_pipeline_parallel_world_size(mpu)

    _ensure_divisibility(world_size, tensor_parallel_size_)
    _ensure_divisibility(dp_world_size, expert_parallel_size_)

    log_dist(
        f"Creating deepspeed groups with model parallel size {tensor_parallel_size_}, "
        f"pipeline parallel size {pp_world_size}, expert parallel size {expert_parallel_size_}, "
        f"world size {world_size}, dp world size {dp_world_size}", [0])

    global _EXPERT_PARALLEL_GROUP, _EXPERT_DATA_PARALLEL_GROUP
    global _EXPERT_PARALLEL_GROUP_RANKS, _EXPERT_DATA_PARALLEL_GROUP_RANKS

    group_name = f"ep_size_{expert_parallel_size_}"

    # Only create groups if they don't already exist
    # Need to check conditions outside the group creation loop because of the way torch.dist group creation works
    if group_name not in _EXPERT_DATA_PARALLEL_GROUP and group_name not in _EXPERT_PARALLEL_GROUP:
        expert_parallel_groups, expert_data_parallel_groups = _get_expert_parallel_ranks(
            world_size, tensor_parallel_size_, expert_parallel_size_, pp_world_size, use_data_before_expert_parallel_)
        for ranks in expert_parallel_groups:
            group = dist.new_group(ranks)
            if rank in list(ranks):
                _EXPERT_PARALLEL_GROUP[group_name] = group
                _EXPERT_PARALLEL_GROUP_RANKS[group_name] = ranks

        for ranks in expert_data_parallel_groups:
            group = dist.new_group(ranks)
            if rank in list(ranks):
                _EXPERT_DATA_PARALLEL_GROUP[group_name] = group
                _EXPERT_DATA_PARALLEL_GROUP_RANKS[group_name] = ranks


def _get_max_expert_size():