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

Code before:
    import apex
    if hasattr(apex, 'optimizers') and hasattr(apex.optimizers, 'FusedAdam'):
        ZERO_SUPPORTED_OPTIMIZERS.append(apex.optimizers.FusedAdam)
except ImportError:
    pass


def is_zero_supported_optimizer(optimizer):
    if dist.get_rank() == 0:
        logger.info(f'Checking ZeRO support for optimizer={optimizer.__class__.__name__} type={type(optimizer)}')
    return type(optimizer) in ZERO_SUPPORTED_OPTIMIZERS


@instrument_w_nvtx
def assert_lst_len_same_as_other_ranks(lst: List[int]) -> None:
    rank0_len_tensor = torch.tensor(
        len(lst) if dist.get_rank() == 0 else -1,
        dtype=int,
        device=torch.device(get_accelerator().device_name(os.environ["LOCAL_RANK"])),
        requires_grad=False,
    )
    local_len_tensor = len(lst)
    dist.broadcast(rank0_len_tensor, src=0, async_op=False)
    rank0_len_tensor_cpu = rank0_len_tensor.cpu().item()
    if rank0_len_tensor_cpu != local_len_tensor:
        raise RuntimeError(f"disagreement between rank0 and rank{dist.get_rank()}: "
                           f"\n rank0: {rank0_len_tensor_cpu} "
                           f"\n rank{dist.get_rank()}: {local_len_tensor}")


def get_lst_from_rank0(lst: List[int]) -> None:
    """
    NOTE: creates both communication and synchronization overhead so should be used
    sparingly
    """
    lst_tensor = torch.tensor(
        lst if dist.get_rank() == 0 else [-1] * len(lst),
        dtype=int,
        device=torch.device(get_accelerator().device_name(os.environ["LOCAL_RANK"])),
        requires_grad=False,
    )
    dist.broadcast(lst_tensor, src=0, async_op=False)

    return [t.item() for t in lst_tensor.cpu()]


@instrument_w_nvtx
def assert_ints_same_as_other_ranks(ints: List[int]) -> None:
    """
    NOTE: creates both communication and synchronization overhead so should be
    used sparingly

Review comment (suggestion):
```suggestion
        raise RuntimeError(f"Detected a disagreement between rank0 and rank{dist.get_rank()}: "
```

Code after:
    import apex
    if hasattr(apex, 'optimizers') and hasattr(apex.optimizers, 'FusedAdam'):
        ZERO_SUPPORTED_OPTIMIZERS.append(apex.optimizers.FusedAdam)
except ImportError:
    pass


def is_zero_supported_optimizer(optimizer):
    if dist.get_rank() == 0:
        logger.info(f'Checking ZeRO support for optimizer={optimizer.__class__.__name__} type={type(optimizer)}')
    return type(optimizer) in ZERO_SUPPORTED_OPTIMIZERS


@instrument_w_nvtx
def assert_lst_len_same_as_other_ranks(lst: List[int]) -> None:
    rank0_len_tensor = torch.tensor(
        len(lst) if dist.get_rank() == 0 else -1,
        dtype=int,
        device=torch.device(get_accelerator().device_name(os.environ["LOCAL_RANK"])),
        requires_grad=False,
    )
    local_list_length = len(lst)
    dist.broadcast(rank0_len_tensor, src=0, async_op=False)
    rank0_list_length = rank0_len_tensor.cpu().item()
    if rank0_list_length != local_list_length:
        raise RuntimeError(f"Detected a disagreement on list length between rank0 and rank{dist.get_rank()}: "
                           f"\n rank0: {rank0_list_length} "
                           f"\n rank{dist.get_rank()}: {local_list_length}")


def get_lst_from_rank0(lst: List[int]) -> None:
    """
    NOTE: creates both communication and synchronization overhead so should be used
    sparingly
    """
    lst_tensor = torch.tensor(
        lst if dist.get_rank() == 0 else [-1] * len(lst),
        dtype=int,
        device=torch.device(get_accelerator().device_name(os.environ["LOCAL_RANK"])),
        requires_grad=False,
    )
    dist.broadcast(lst_tensor, src=0, async_op=False)

    return [t.item() for t in lst_tensor.cpu()]


@instrument_w_nvtx
def assert_ints_same_as_other_ranks(ints: List[int]) -> None:
    """
    NOTE: creates both communication and synchronization overhead so should be
    used sparingly