p2p.py

import dataclasses
from typing import List, Sequence, Tuple

import torch
from nanotron import distributed as dist
from nanotron import logging
from nanotron.utils import get_untyped_storage, tensor_from_untyped_storage

logger = logging.get_logger(__name__)

FIRST_METADATA_SIZE = 7
SECOND_METADATA_SIZE = 1024
ID_TO_DTYPE = [
    torch.float32,
    torch.float64,
    torch.complex64,
    torch.complex128,
    torch.float16,
    torch.bfloat16,
    torch.uint8,
    torch.int8,
    torch.int16,
    torch.int32,
    torch.int64,
    torch.bool,
]
DTYPE_TO_ID = {dtype: id_ for id_, dtype in enumerate(ID_TO_DTYPE)}

ID_TO_REQUIRES_GRAD = [True, False]
REQUIRES_GRAD_TO_ID = {value: id_ for id_, value in enumerate(ID_TO_REQUIRES_GRAD)}
ID_TO_IS_CONTIGUOUS = [True, False]
IS_CONTIGUOUS_TO_ID = {value: id_ for id_, value in enumerate(ID_TO_IS_CONTIGUOUS)}


@dataclasses.dataclass
class P2PTensorMetaData:
    shape: Sequence[int]
    stride: Sequence[int]
    is_contiguous: bool
    untyped_storage_size: int
    storage_offset: int
    dtype: torch.dtype
    requires_grad: bool

    def create_empty_storage(self, device: torch.device) -> torch.Tensor:
        buffer = torch.empty(
            size=(self.untyped_storage_size,),
            requires_grad=False,
            dtype=torch.int8,
            device=device,
            memory_format=torch.contiguous_format,
        ).view(dtype=self.dtype)
        buffer.requires_grad = self.requires_grad

        if self.is_contiguous:
            buffer = buffer.as_strided(
                size=tuple(self.shape), stride=tuple(self.stride), storage_offset=self.storage_offset
            )

        # Complex needs to be viewed as real first
        # TODO @thomasw21: Find the issue with send/recv complex tensors
        buffer = torch.view_as_real(buffer) if self.dtype.is_complex else buffer

        return buffer

    def reshape(self, buffer):
        """Changes the way we view buffer in order to fit metadata"""
        # TODO @thomasw21: Find the issue with send/recv complex tensors
        buffer = torch.view_as_complex(buffer) if self.dtype.is_complex else buffer

        # Set shape and stride
        if not self.is_contiguous:
            buffer = buffer.as_strided(
                size=tuple(self.shape), stride=tuple(self.stride), storage_offset=self.storage_offset
            )

        return buffer

    @staticmethod
    def to_first_metadata(tensor: torch.Tensor, device: torch.device) -> torch.Tensor:
        # TODO @nouamane: avoid having two metadata comms, and preallocate shape/stride instead
        return torch.tensor(
            [
                len(tensor.shape),
                len(tensor.stride()),
                IS_CONTIGUOUS_TO_ID[tensor.is_contiguous()],
                get_untyped_storage(tensor).size(),
                tensor.storage_offset(),
                DTYPE_TO_ID[tensor.dtype],
                REQUIRES_GRAD_TO_ID[tensor.requires_grad],
            ],
            dtype=torch.long,
            device=device,
        )

    @staticmethod
    def to_second_metadata(tensor: torch.Tensor, device: torch.device) -> torch.Tensor:
        return torch.tensor(tensor.shape + tensor.stride(), dtype=torch.long, device=device)

    @classmethod
    def from_metadata(cls, first_metadata: List[int], second_metadata: List[int]):
        shape_and_stride = second_metadata
        (
            num_shape,
            num_stride,
            is_contiguous,
            untyped_storage_size,
            storage_offset,
            dtype_id,
            requires_grad_id,
        ) = first_metadata
        return cls(
            shape=shape_and_stride[: len(shape_and_stride) // 2],
            stride=shape_and_stride[len(shape_and_stride) // 2 :],
            is_contiguous=ID_TO_IS_CONTIGUOUS[is_contiguous],
            untyped_storage_size=untyped_storage_size,
            storage_offset=storage_offset,
            dtype=ID_TO_DTYPE[dtype_id],
            requires_grad=ID_TO_REQUIRES_GRAD[requires_grad_id],
        )


def view_as_contiguous(tensor: torch.Tensor):
    """Given a tensor, we want to view the tensor as a contiguous storage"""
    tensor_numel = tensor.numel()
    tensor_element_size = tensor.element_size()
    untyped_storage = get_untyped_storage(tensor)
    untyped_storage_size = untyped_storage.size()
    untyped_element_size = untyped_storage.element_size()
    assert (
        tensor_numel * tensor_element_size >= untyped_storage_size * untyped_element_size
    ), "Expect storage_size to be smaller than tensor size. It might not be true, when you use slicing for example though. We probably don't want to support it in our P2P system"
    buffer = tensor_from_untyped_storage(untyped_storage=untyped_storage, dtype=tensor.dtype)
    return buffer


class P2P:
    def __init__(self, pg: dist.ProcessGroup, device: torch.device):
        self.pg = pg
        self.device = device
        self.first_metadata = torch.empty(FIRST_METADATA_SIZE, dtype=torch.long, device=self.device)
        self.second_metadata = torch.empty(SECOND_METADATA_SIZE, dtype=torch.long, device=self.device)

    def _send_first_metadata_p2p_op(self, tensor: torch.Tensor, to_rank: int, tag: int = 0) -> dist.P2POp:
        first_metadata = P2PTensorMetaData.to_first_metadata(tensor=tensor, device=self.device)
        return dist.P2POp(
            op=dist.isend,
            tensor=first_metadata,
            peer=dist.get_global_rank(group=self.pg, group_rank=to_rank),
            group=self.pg,
            tag=tag,
        )

    def _recv_first_metadata_p2p_op(self, from_rank: int, tag: int = 0) -> Tuple[torch.Tensor, dist.P2POp]:
        first_metadata_buffer = torch.empty((FIRST_METADATA_SIZE,), dtype=torch.long, device=self.device)
        return first_metadata_buffer, dist.P2POp(
            op=dist.irecv,
            tensor=first_metadata_buffer,
            peer=dist.get_global_rank(group=self.pg, group_rank=from_rank),
            group=self.pg,
            tag=tag,
        )

    def _send_second_metadata_p2p_op(self, tensor: torch.Tensor, to_rank: int, tag: int = 0) -> dist.P2POp:
        second_metadata = P2PTensorMetaData.to_second_metadata(tensor=tensor, device=self.device)
        return dist.P2POp(
            op=dist.isend,
            tensor=second_metadata,
            peer=dist.get_global_rank(group=self.pg, group_rank=to_rank),
            group=self.pg,
            tag=tag,
        )

    def _recv_second_metadata_p2p_op(
        self, shape_length: int, stride_length: int, from_rank: int, tag: int = 0
    ) -> Tuple[torch.Tensor, dist.P2POp]:
        second_metadata_buffer = torch.empty((shape_length + stride_length,), dtype=torch.long, device=self.device)
        return second_metadata_buffer, dist.P2POp(
            op=dist.irecv,
            tensor=second_metadata_buffer,
            peer=dist.get_global_rank(group=self.pg, group_rank=from_rank),
            group=self.pg,
            tag=tag,
        )

    def _send_data_p2p_op(self, tensor: torch.Tensor, to_rank: int, tag: int = 0) -> dist.P2POp:
        return dist.P2POp(
            op=dist.isend,
            tensor=tensor,
            peer=dist.get_global_rank(group=self.pg, group_rank=to_rank),
            group=self.pg,
            tag=tag,
        )

    def _recv_data_p2p_op(
        self, tensor_metadata: P2PTensorMetaData, from_rank: int, tag: int = 0
    ) -> Tuple[torch.Tensor, dist.P2POp]:
        tensor_buffer = tensor_metadata.create_empty_storage(self.device)
        return tensor_buffer, dist.P2POp(
            op=dist.irecv,
            tensor=tensor_buffer,
            peer=dist.get_global_rank(group=self.pg, group_rank=from_rank),
            group=self.pg,
            tag=tag,
        )

    def _send_meta(self, tensor: torch.Tensor, to_rank: int, tag: int):
        cpu_tensor = torch.tensor(
            [
                len(tensor.shape),
                len(tensor.stride()),
                IS_CONTIGUOUS_TO_ID[tensor.is_contiguous()],
                get_untyped_storage(tensor).size(),
                tensor.storage_offset(),
                DTYPE_TO_ID[tensor.dtype],
                REQUIRES_GRAD_TO_ID[tensor.requires_grad],
            ],
            dtype=torch.long,
        )
        self.first_metadata.copy_(cpu_tensor)
        dist.send(
            self.first_metadata,
            dst=dist.get_global_rank(group=self.pg, group_rank=to_rank),
            group=self.pg,
            tag=tag,
        )

        second_metadata = tensor.shape + tensor.stride()
        assert len(tensor.shape) == self.first_metadata[0]
        assert len(tensor.stride()) == self.first_metadata[1]

        # increase buffer size
        if len(second_metadata) > len(self.second_metadata):
            self.second_metadata = torch.empty(len(second_metadata), dtype=torch.long, device=self.device)

        self.second_metadata[: len(second_metadata)].copy_(torch.tensor(second_metadata, dtype=torch.long))

        dist.send(
            self.second_metadata[: len(second_metadata)],
            dst=dist.get_global_rank(group=self.pg, group_rank=to_rank),
            group=self.pg,
            tag=tag,
        )

    def _recv_meta(self, from_rank: int, tag: int) -> P2PTensorMetaData:
        dist.recv(
            self.first_metadata,
            src=dist.get_global_rank(group=self.pg, group_rank=from_rank),
            group=self.pg,
            tag=tag,
        )
        (
            num_shape,
            num_stride,
            is_contiguous,
            untyped_storage_size,
            storage_offset,
            dtype_id,
            requires_grad_id,
        ) = self.first_metadata

        # self.pg.recv([second], from_rank, 0).wait() # more direct API
        second_metadata_num_elements = num_shape + num_stride

        # increase buffer size
        if second_metadata_num_elements > len(self.second_metadata):
            self.second_metadata = torch.empty(second_metadata_num_elements, dtype=torch.long, device=self.device)

        dist.recv(
            self.second_metadata[:second_metadata_num_elements],
            src=dist.get_global_rank(group=self.pg, group_rank=from_rank),
            group=self.pg,
            tag=tag,
        )

        shape = self.second_metadata[:num_shape]
        stride = self.second_metadata[num_shape:second_metadata_num_elements]

        return P2PTensorMetaData(
            dtype=ID_TO_DTYPE[dtype_id],
            requires_grad=ID_TO_REQUIRES_GRAD[requires_grad_id],
            shape=shape,
            stride=stride,
            is_contiguous=ID_TO_IS_CONTIGUOUS[is_contiguous],
            untyped_storage_size=untyped_storage_size,
            storage_offset=storage_offset,
        )

    def isend_tensors(self, tensors: List[torch.Tensor], to_rank: int, tag: int = 0) -> List[dist.Work]:
        futures = []
        current_rank = dist.get_rank(self.pg)
        logger.debug(f"Current rank {current_rank} sending to rank {to_rank}. Nb_tensors: {len(tensors)}")
        for tensor in tensors:
            if to_rank != current_rank:
                self._send_meta(tensor, to_rank=to_rank, tag=tag)
                if tensor.is_contiguous():
                    buffer = tensor
                else:
                    # If the tensor is not contiguous we send the entire storage
                    buffer = view_as_contiguous(tensor)

                # TODO @thomasw21: Find the issue with send/recv complex tensors
                buffer = torch.view_as_real(buffer) if buffer.is_complex() else buffer

                futures.append(
                    dist.isend(
                        buffer,
                        dst=dist.get_global_rank(group=self.pg, group_rank=to_rank),
                        group=self.pg,
                        tag=tag,
                    )
                )
            else:
                raise ValueError("Tried sending tensor to itself")
        return futures

    def irecv_tensors(
        self, num_tensors: int, from_rank: int, tag: int = 0
    ) -> Tuple[List[torch.Tensor], List[dist.Work]]:
        futures = []
        buffers = []
        current_rank = dist.get_rank(self.pg)
        logger.debug(f"Current rank {current_rank} receiving from rank {from_rank}. Nb_tensors: {num_tensors}")
        for _ in range(num_tensors):
            if from_rank != current_rank:
                meta = self._recv_meta(from_rank=from_rank, tag=tag)

                buffer = meta.create_empty_storage(device=self.device)

                futures.append(
                    dist.irecv(
                        buffer,
                        src=dist.get_global_rank(group=self.pg, group_rank=from_rank),
                        group=self.pg,
                        tag=tag,
                    )
                )

                buffer = meta.reshape(buffer=buffer)

                # Add to the list
                buffers.append(buffer)
            else:
                raise ValueError("Tried receiving tensor from itself")
        return buffers, futures

    def send_tensors(self, tensors: List[torch.Tensor], to_rank: int, tag: int = 0):
        futures = self.isend_tensors(tensors=tensors, to_rank=to_rank, tag=tag)
        for future in futures:
            future.wait()

    def recv_tensors(self, num_tensors: int, from_rank: int, tag: int = 0) -> List[torch.Tensor]:
        buffers, futures = self.irecv_tensors(num_tensors=num_tensors, from_rank=from_rank, tag=tag)
        for future in futures:
            future.wait()
        return buffers


class BatchTensorSendRecvState:
    """
    This class is used to register send/recv batches of tensors, and
    then executes send/recv in `flush()` calls. This is useful for
    amortizing the cost of sending and receiving tensors over multiple
    iterations.
    """

    p2p: P2P
    first_metadata_p2p_ops: List[dist.P2POp]
    second_metadata_p2p_ops: List[dist.P2POp]
    data_p2p_ops: List[dist.P2POp]
    recv_first_metadata_buffers: List[torch.Tensor]
    recv_from_ranks: List[int]

    def __init__(self, p2p: P2P):
        self.p2p = p2p
        self._reset()

    def _reset(self):
        self.first_metadata_p2p_ops: List[dist.P2POp] = []
        self.second_metadata_p2p_ops: List[dist.P2POp] = []
        self.data_p2p_ops: List[dist.P2POp] = []
        self.recv_first_metadata_buffers: List[torch.Tensor] = []
        self.recv_from_ranks: List[int] = []

    def __str__(self):
        return f"BatchTensorSendRecvState(first_metadata_p2p_ops={len(self.first_metadata_p2p_ops)}, second_metadata_p2p_ops={len(self.second_metadata_p2p_ops)}, data_p2p_ops={len(self.data_p2p_ops)}, recv_first_metadata_buffers={len(self.recv_first_metadata_buffers)}, recv_from_ranks={self.recv_from_ranks})"

    def add_send(self, tensor: torch.Tensor, to_rank: int, tag: int = 0):
        self.first_metadata_p2p_ops.append(
            self.p2p._send_first_metadata_p2p_op(tensor=tensor, to_rank=to_rank, tag=tag)
        )
        self.second_metadata_p2p_ops.append(
            self.p2p._send_second_metadata_p2p_op(tensor=tensor, to_rank=to_rank, tag=tag)
        )
        self.data_p2p_ops.append(
            self.p2p._send_data_p2p_op(tensor=view_as_contiguous(tensor), to_rank=to_rank, tag=tag)
        )

    def add_recv(self, from_rank: int, tag: int = 0) -> int:
        """
        Only add p2p ops for the first operation, as `_recv_second_metadata` and `_recv_data_p2p_op`
        require results from the first metadata to be transfered first.
        Return: index of the recv_buffer in `self.recv_first_metadata_buffers`
        """
        buffer, recv_op = self.p2p._recv_first_metadata_p2p_op(from_rank=from_rank, tag=tag)
        self.first_metadata_p2p_ops.append(recv_op)
        self.recv_first_metadata_buffers.append(buffer)
        self.recv_from_ranks.append(from_rank)
        return len(self.recv_first_metadata_buffers) - 1

    def _send_recv_first_metadata(self) -> List[List[int]]:
        # Send/Recv first metadata
        reqs = dist.batch_isend_irecv(self.first_metadata_p2p_ops)
        for req in reqs:
            req.wait()
        # We want an early cpu/gpu sync here as we are right after the wait so it's nearly free.
        # Removing the tolist call here delays the sync and will impact performance.
        # We need to instantiate it in a list because it is used twice
        first_metadatas = [tensor.tolist() for tensor in self.recv_first_metadata_buffers]
        return first_metadatas

    def _send_recv_second_metadata(self, first_metadata: List[List[int]]) -> List[List[int]]:
        # turn a list of tuple into a tuple of list
        recv_second_metadata_buffers, recv_second_metadata_ops = zip(
            *(
                self.p2p._recv_second_metadata_p2p_op(
                    shape_length=num_shape, stride_length=num_stride, from_rank=from_rank
                )
                for (num_shape, num_stride, *_), from_rank in zip(first_metadata, self.recv_from_ranks)
            )
        )
        recv_second_metadata_ops = list(recv_second_metadata_ops)
        # Send/Recv second metadata
        reqs = dist.batch_isend_irecv(self.second_metadata_p2p_ops + recv_second_metadata_ops)
        for req in reqs:
            req.wait()

        # We want an early cpu/gpu sync here as we are right after the wait so it's nearly free.
        # Removing the tolist call here delays the sync and will impact performance.
        second_metadatas = [tensor.tolist() for tensor in recv_second_metadata_buffers]
        return second_metadatas

    def _send_recv_data(self, tensor_metadatas: List[P2PTensorMetaData]) -> List[torch.Tensor]:
        # turn a list of tuples into a tuple of list
        recv_data_buffers, recv_data_ops = zip(
            *(
                self.p2p._recv_data_p2p_op(tensor_metadata=tensor_metadata, from_rank=from_rank)
                for tensor_metadata, from_rank in zip(tensor_metadatas, self.recv_from_ranks)
            )
        )
        recv_data_ops = list(recv_data_ops)
        # Send/Recv tensor data
        futures = dist.batch_isend_irecv(self.data_p2p_ops + recv_data_ops)
        for future in futures:
            future.wait()

        # Format tensor by setting the stride
        return [
            recv_data_buffer.as_strided(size=tuple(tensor_metadata.shape), stride=tuple(tensor_metadata.stride))
            for recv_data_buffer, tensor_metadata in zip(recv_data_buffers, tensor_metadatas)
        ]

    def flush(self) -> List[torch.Tensor]:
        """
        Run all communication in a batch.
        Return `torch.Tensor` in the case of recv.
        """
        assert len(self.recv_first_metadata_buffers) == len(
            self.recv_from_ranks
        ), f"len(self.recv_first_metadata_buffers)={len(self.recv_first_metadata_buffers)}, len(self.recv_from_ranks)={len(self.recv_from_ranks)} but should be equal."

        # If there is no communication, return
        if len(self.first_metadata_p2p_ops) == 0:
            return []

        # If there is no recv
        if len(self.recv_first_metadata_buffers) == 0:
            reqs = dist.batch_isend_irecv(
                self.first_metadata_p2p_ops + self.second_metadata_p2p_ops + self.data_p2p_ops
            )
            for req in reqs:
                req.wait()
            self._reset()
            return []

        # Send/Recv first metadata
        logger.debug(f"First metadata: {[p2pop.op for p2pop in self.first_metadata_p2p_ops]}")
        # TODO(kunhao): We could actually send all at once like the above no recv case. But I need to benchmark the performance.
        first_metadatas = self._send_recv_first_metadata()
        # Send/Recv second metadata
        second_metadatas = self._send_recv_second_metadata(first_metadatas)

        tensor_metadatas = [
            P2PTensorMetaData.from_metadata(first_metadata, second_metadata)
            for first_metadata, second_metadata in zip(first_metadatas, second_metadatas)
        ]

        recv_tensors = self._send_recv_data(tensor_metadatas)
        # Reset state
        self._reset()

        return recv_tensors