Repository: hpcaitech/ColossalAI
File: colossalai/pipeline/p2p.py (Python)

Code before:
    send_group: Optional[ProcessGroup] = None,
    recv_group: Optional[ProcessGroup] = None,
) -> Any:
    if c10d._rank_not_in_group(send_group) or c10d._rank_not_in_group(recv_group):
        c10d._warn_not_in_group("_communicate")
        return

    current_send_device, is_send_nccl_backend = check_device(send_group)
    current_recv_device, is_recv_nccl_backend = check_device(recv_group)

    is_nccl_backend = is_send_nccl_backend and is_recv_nccl_backend

    assert current_send_device == current_recv_device
    current_device = current_send_device

    assert (send_dst is not None) or (recv_src is not None)

    can_fast_send = False
    send_metadata = None
    if send_dst is not None:
        can_fast_send = _check_if_fast_send_available(object) and is_nccl_backend
        if not can_fast_send:
            send_metadata = ['__serialization__', object]
        else:
            send_metadata = []
            send_metadata.append('__tensor__')
            if type(object) is torch.Tensor:
                send_metadata.append(0)
                send_metadata.append((object.shape, object.dtype, object.requires_grad))
            elif type(object) is list:
                send_metadata.append(1)
                for v in object:
                    send_metadata.append((v.shape, v.dtype, v.requires_grad))
            elif type(object) is dict:
                send_metadata.append(2)
                for k, v in object.items():
                    send_metadata.append((k, v.shape, v.dtype, v.requires_grad))

    recv_metadata = _send_recv_serialization_object(send_metadata, send_dst, recv_src, send_group, recv_group, current_device, is_nccl_backend)
    if recv_metadata is not None:
        assert type(recv_metadata) is list and len(recv_metadata) >= 2
        if recv_metadata[0] == '__serialization__':
            return recv_metadata[1]
        else:
            recv_metadata.pop(0)
    if not can_fast_send and send_dst is not None:
        return

    send_tensor_list = None
    if type(object) is torch.Tensor:
        send_tensor_list = object

Review comment (suggestion):
It seems `send_metadata` is always serialized before sending. Then why not use a dataclass to represent it? This would be more readable.

Code after:

    if recv_object_tensor is not None and recv_object_size_tensor is not None:
        recv_object_tensor = recv_object_tensor.type(torch.uint8)
        if recv_object_tensor.device != torch.device("cpu"):
            recv_object_tensor = recv_object_tensor.cpu()

        unpickle_object = _cuda_safe_tensor_to_object(
            recv_object_tensor, recv_object_size_tensor.item())

        if (
            isinstance(unpickle_object, torch.Tensor)
            and unpickle_object.device.index != torch.cuda.current_device()
        ):
            unpickle_object = unpickle_object.cuda()

        return unpickle_object


def _check_if_fast_send_available(object):
    if type(object) is torch.Tensor:
        return True
    elif type(object) is list:
        is_list_of_tensor = all([type(v) is torch.Tensor for v in object])
        return is_list_of_tensor
    elif type(object) is dict:
        is_dict_of_tensor = all([type(k) is str and type(
            v) is torch.Tensor for k, v in object.items()])

        return is_dict_of_tensor
    return False


def _communicate(
    object,
    send_dst: Optional[int],
    recv_src: Optional[int],
    send_group: Optional[ProcessGroup] = None,
    recv_group: Optional[ProcessGroup] = None,
) -> Any:
    if c10d._rank_not_in_group(send_group) or c10d._rank_not_in_group(recv_group):
        c10d._warn_not_in_group("_communicate")
        return

    current_send_device, is_send_nccl_backend = check_device(send_group)
    current_recv_device, is_recv_nccl_backend = check_device(recv_group)

    is_nccl_backend = is_send_nccl_backend and is_recv_nccl_backend

    assert current_send_device == current_recv_device
    current_device = current_send_device