Function process_group_is_intranode

colossalai/quantization/fp8.py:33–48

(pg)

def process_group_is_intranode(pg):
    if pg is None:
        from torch.distributed.distributed_c10d import _get_default_group

        pg = _get_default_group()

    local_world_size = None
    for var in ["LOCAL_WORLD_SIZE", "OMPI_COMM_WORLD_LOCAL_SIZE", "SLURM_TASKS_PER_NODE"]:
        if var in os.environ:
            local_world_size = int(os.environ["LOCAL_WORLD_SIZE"])
    if local_world_size is None:
        local_world_size = torch.cuda.device_count()

    group_ranks = dist.get_process_group_ranks(pg)
    group_ranks_node_ids = [rank // local_world_size for rank in group_ranks]
    return min(group_ranks_node_ids) == max(group_ranks_node_ids)


def cast_to_fp8(