Function _create_expert_and_data_parallel

deepspeed/utils/groups.py:240–382

Create expert and data parallel groups. When mp_size is None or 1: legacy consecutive ordering (backward compatible). When mp_size > 1 and mp_mode=="tp": TP-strided rank ordering. When mp_size > 1 and mp_mode=="sp": consecutive rank ordering. Note: Caller of this function is respon

(expert_parallel_size_,
                                     mp_size=None,
                                     pp_size=None,
                                     mp_mode="tp",
                                     use_data_before_expert_parallel_=False)

def _create_expert_and_data_parallel(expert_parallel_size_,
                                     mp_size=None,
                                     pp_size=None,
                                     mp_mode="tp",
                                     use_data_before_expert_parallel_=False):
    """Create expert and data parallel groups.

    When mp_size is None or 1: legacy consecutive ordering (backward compatible).
    When mp_size > 1 and mp_mode=="tp": TP-strided rank ordering.
    When mp_size > 1 and mp_mode=="sp": consecutive rank ordering.

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

    Example - E + D parallel (legacy path)
    world_size = 16
    expert_parallel_size = 2 # number of experts in same group
    expert_data_parallel_group = [0,2,4,6,8,10,12,14], [1,3,5,7,9,11,13,15] - all reduce is only on MoE params
    expert_parallel_group = [0, 1], [2,3], [4,5], [6,7], [8,9] - no all reduce, but all to all
    data_parallel_group = [0,1,...,15] - all reduce is only on non-MoE

    Args:
        expert_parallel_size_ (int): Expert parallel group size.
        mp_size (int, optional): Model parallel size (TP or SP). None treated as 1.
        pp_size (int, optional): Pipeline parallel size. None falls back to mpu.
        mp_mode (str): "tp" for TP-strided ordering, "sp" for consecutive ordering.
        use_data_before_expert_parallel_ (bool): Use the D + E instead of E + D topology.
    """
    assert dist.is_initialized()

    # Resolve parameters for backward compat
    effective_mp_size = 1 if mp_size is None else mp_size

    log_dist(f'Creating expert and data parallel groups with size {expert_parallel_size_}', ranks=[0])
    world_size = dist.get_world_size()

    # Resolve pp_size
    if pp_size is not None:
        pp_world_size = pp_size
    else:
        pp_world_size = 1 if mpu is None else bwc_pipeline_parallel_world_size(mpu)

    rank = dist.get_rank()

    pp_stride = world_size // pp_world_size
    _ensure_divisibility(pp_stride, expert_parallel_size_)

    group_name = f"ep_size_{expert_parallel_size_}"

    global _EXPERT_DATA_PARALLEL_GROUP
    global _EXPERT_DATA_PARALLEL_GROUP_RANKS
    global _EXPERT_PARALLEL_GROUP
    global _EXPERT_PARALLEL_GROUP_RANKS

    # Legacy path: mp_size <= 1 (preserves exact original behavior)
    if effective_mp_size <= 1:
        ep_stride = pp_stride // expert_parallel_size_

        # Build the expert data parallel groups.