Repository: microsoft/DeepSpeed
File: deepspeed/utils/groups.py (Python)

Code before:
def _get_sequence_parallel_world_size():
    """Return world size for the model parallel group."""
    global mpu
    if mesh_device is not None:
        return dist.get_world_size(mesh_device.get_group(mesh_dim="sequence_parallel"))
    if mpu is not None and hasattr(mpu, 'get_sequence_parallel_world_size'):
        return mpu.get_sequence_parallel_world_size()
    return 1


def _get_sequence_parallel_rank():
    """Return my rank for the data parallel group."""
    global mpu
    if mpu is not None and hasattr(mpu, 'get_sequence_parallel_rank'):
        return mpu.get_sequence_parallel_rank()
    return 0


def _get_sequence_parallel_group():
    global mpu
    if mpu is not None and hasattr(mpu, 'get_sequence_parallel_group'):
        return mpu.get_sequence_parallel_group()
    elif mesh_device is not None:
        return mesh_device.get_group(mesh_dim="sequence_parallel")
    else:
        KeyError("No sequence parallel group found")
        ##Should this be None?
    #return None


def _get_sequence_data_parallel_world_size():
    """Return world size for the model parallel group."""
    global mpu
    if mpu is not None and hasattr(mpu, 'get_sequence_data_parallel_world_size'):
        return mpu.get_sequence_data_parallel_world_size()
    return _get_data_parallel_world_size()


def _get_sequence_data_parallel_rank():
    """Return my rank for the data parallel group."""
    global mpu
    if mpu is not None and hasattr(mpu, 'get_sequence_data_parallel_rank'):
        return mpu.get_sequence_data_parallel_rank()
    return _get_data_parallel_rank()


def _get_sequence_data_parallel_group():
    global mpu
    # When sequence parallelism is enabled, the process group for zero sharding and
    # gradient allreduce must be across both dimensions of data and sequence parallelism.
    if mpu is not None and hasattr(mpu, 'get_sequence_data_parallel_group'):

Review comment (bug):
Failing fast is preferred, so raising an exception is the right thing to do here (asking for the SP group when SP is not enabled should not happen).

Code after:
    global mpu
    if mesh_device is not None:
        return dist.get_world_size(mesh_device.get_group(mesh_dim="sequence_parallel"))
    if mpu is not None and hasattr(mpu, 'get_sequence_parallel_world_size'):
        return mpu.get_sequence_parallel_world_size()
    return 1


def _get_sequence_parallel_rank():
    """Return my rank for the data parallel group."""
    global mpu
    if mpu is not None and hasattr(mpu, 'get_sequence_parallel_rank'):
        return mpu.get_sequence_parallel_rank()
    return 0


def _get_sequence_parallel_group():
    global mpu
    if mpu is None or not hasattr(mpu, 'get_sequence_parallel_group'):
        if mesh_device is None:
            raise KeyError("No sequence parallel group found")
        return mesh_device.get_group(mesh_dim="sequence_parallel")
    return mpu.get_sequence_parallel_group()


def _get_sequence_data_parallel_world_size():
    """Return world size for the model parallel group."""
    global mpu
    if mpu is not None and hasattr(mpu, 'get_sequence_data_parallel_world_size'):
        return mpu.get_sequence_data_parallel_world_size()
    return _get_data_parallel_world_size()


def _get_sequence_data_parallel_rank():
    """Return my rank for the data parallel group."""
    global mpu
    if mpu is not None and hasattr(mpu, 'get_sequence_data_parallel_rank'):
        return mpu.get_sequence_data_parallel_rank()
    return _get_data_parallel_rank()


def _get_sequence_data_parallel_group():
    global mpu
    # When sequence parallelism is enabled, the process group for zero sharding and
    # gradient allreduce must be across both dimensions of data and sequence parallelism.
    if mpu is not None and hasattr(mpu, 'get_sequence_data_parallel_group'):
        return mpu.get_sequence_data_parallel_group()
    return _get_data_parallel_group()


def _get_expert_model_parallel_world_size():