Function get_dist_info

modelscope/utils/torch_utils.py:124–143

Get dist info of a specified group Args: group: The parallel group, default None, for the global group Returns: A tuple of the current rank and world_size of the group

(group=None)

def get_dist_info(group=None) -> Tuple[int, int]:
    """Get dist info of a specified group

    Args:
        group: The parallel group, default None, for the global group

    Returns:
        A tuple of the current rank and world_size of the group
    """
    if is_dist():
        from modelscope.utils.megatron_utils import is_megatron_initialized
        if group is None and is_megatron_initialized():
            from megatron_util import mpu
            group = mpu.get_data_parallel_group()
        rank = dist.get_rank(group)
        world_size = dist.get_world_size(group)
    else:
        rank = 0
        world_size = 1
    return rank, world_size


def get_local_rank():