hivemind/hivemind/averaging/allreduce.py

import asyncio
from enum import Enum
from typing import Any, AsyncIterator, Dict, Optional, Sequence, Tuple, Type

import torch

from hivemind.averaging.partition import AllreduceException, TensorPartContainer, TensorPartReducer
from hivemind.compression import deserialize_torch_tensor, serialize_torch_tensor
from hivemind.p2p import P2P, P2PContext, PeerID, ServicerBase, StubBase
from hivemind.proto import averaging_pb2
from hivemind.utils import get_logger
from hivemind.utils.asyncio import achain, aenumerate, afirst, amap_in_executor, anext, as_aiter

# flavour types
GroupID = bytes
logger = get_logger(__name__)


class AveragingMode(Enum):
    NODE = 0
    CLIENT = 1
    AUX = 2


class AllReduceRunner(ServicerBase):
    """
    An internal class that runs butterfly AllReduce in a predefined group of averagers.

    This class inherits hivemind.p2p.ServicerBase, so it can be used as an RPCServicer for testing purposes without
    creating a full DecentralizedAverager.

    :note: this class returns **differences** between averaged and local tensors in order to improve numerical stability
    :param p2p: a hivemind.p2p.P2P instance used for communication with other peers
    :param servicer_type: a hivemind.p2p.ServicerBase subclass whose RPC signatures are used
      when requesting other peers. Typically, it is DecentralizedAverager, its derivative,
      or AllReduceRunner itself (for testing purposes).
    :param prefix: namespace for servicer's RPCs (typically, equal to prefix for group keys)
    :param group_id: unique identifier of this specific all-reduce run
    :param tensors: local tensors that should be averaged with groupmates
    :param tensors: local tensors that should be averaged with groupmates
    :param peer_id: your peer_id, must be included in ordered_peer_ids
    :param ordered_peer_ids: group peer_ids ordered s.t. i-th peer_id is responsible for averaging i-th part
    :param peer_fractions: for each peer, a target fraction of vector elements that this peer should average
      (the actual number of values by peer will be nearly proportional, but there are no exact guarantees)
    :param modes: AveragingMode for each peer in ordered_peer_ids (normal, client-only or auxiliary)
    :param weights: scaling coefficients for weighted averaging (default = equal weights for all non-aux peers)
    :param gathered: additional user-defined data collected from this group
    :param kwargs: additional paramters (e.g. part_size_bytes) will be passed to TensorPartContainer
    """

    def __init__(
        self,
        *,
        p2p: P2P,
        servicer_type: Type[ServicerBase],
        prefix: Optional[str],
        group_id: GroupID,
        tensors: Sequence[torch.Tensor],
        ordered_peer_ids: Sequence[PeerID],
        peer_fractions: Tuple[float, ...],
        weights: Optional[Sequence[float]] = None,
        modes: Optional[Sequence[AveragingMode]] = None,
        gathered: Optional[Dict[PeerID, Any]] = None,
        **kwargs,
    ):
        self._p2p = p2p
        self.peer_id = p2p.peer_id
        assert self.peer_id in ordered_peer_ids, "peer_id is not a part of the group"

        if not issubclass(servicer_type, ServicerBase):
            raise TypeError("`servicer_type` is expected to be a ServicerBase subclass")
        self._servicer_type = servicer_type
        self._prefix = prefix

        modes = modes or tuple(AveragingMode.CLIENT if frac == 0 else AveragingMode.NODE for frac in peer_fractions)
        weights = weights or tuple(int(mode != AveragingMode.AUX) for mode in modes)
        assert len(weights) == len(modes) == len(ordered_peer_ids), "lists have inconsistent length"
        assert any(mode != AveragingMode.CLIENT for mode in modes), "cannot run allreduce without reducers"
        for mode, frac, weight in zip(modes, peer_fractions, weights):
            assert mode != AveragingMode.CLIENT or frac == 0, "client-mode peer should have zero all-reduce fraction"
            assert mode != AveragingMode.AUX or weight == 0, "auxiliary peer should have zero averaging weight"

        self.group_id, self.ordered_peer_ids = group_id, ordered_peer_ids
        self.modes, self.peer_fractions, self.gathered = modes, peer_fractions, gathered

        self._future = asyncio.Future()

        self.sender_peer_ids, self.sender_weights = [], []
        for peer_id, weight, mode in zip(self.ordered_peer_ids, weights, modes):
            if mode != AveragingMode.AUX:
                self.sender_peer_ids.append(peer_id)
                self.sender_weights.append(weight)

        peer_id_index = self.ordered_peer_ids.index(self.peer_id)
        self.tensor_part_container = TensorPartContainer(tensors, peer_fractions, **kwargs)
        self.parts_for_local_averaging = self.tensor_part_container.get_raw_input_parts(peer_id_index)
        self.tensor_part_reducer = TensorPartReducer(
            tuple(part.shape for part in self.parts_for_local_averaging),
            len(self.sender_peer_ids),
            self.sender_weights,
        )

    def __repr__(self):
        return f"{self.__class__.__name__}({self.peer_id}, group_size={self.group_size})"

    def __aiter__(self):
        return self.run()

    def __contains__(self, peer_id: PeerID):
        return peer_id in self.ordered_peer_ids

    @property
    def group_size(self):
        return len(self.ordered_peer_ids)

    def _get_peer_stub(self, peer: PeerID) -> StubBase:
        return self._servicer_type.get_stub(self._p2p, peer, namespace=self._prefix)

    async def run(self) -> AsyncIterator[torch.Tensor]:
        """Run all-reduce, return differences between averaged and original tensors as they are computed"""
        pending_tasks = set()
        try:
            if len(self.sender_peer_ids) == 0:
                logger.debug(f"{self} - finished all-reduce early: all peers are auxiliaries ({self.modes})")
                self.finalize()

            elif self.peer_id in self.sender_peer_ids:
                for peer_id, parts in zip(self.ordered_peer_ids, self.tensor_part_container.num_parts_by_peer):
                    if parts != 0:
                        pending_tasks.add(asyncio.create_task(self._communicate_with_peer(peer_id)))

                async for averaged_tensor_delta in self.tensor_part_container.iterate_output_tensors():
                    yield averaged_tensor_delta  # delta = averaged_tensor - original_tensor
                self.finalize()

            else:  # auxiliary peer
                await self.tensor_part_reducer.finished.wait()
                self.finalize()

        except BaseException as e:
            self.finalize(exception=e)
            for task in pending_tasks:
                task.cancel()
            raise

    async def _communicate_with_peer(self, peer_id: PeerID):
        """Send a part of local tensors and metadata to a single peer, receive the average for that part of tensors"""
        peer_index = self.ordered_peer_ids.index(peer_id)
        if peer_id == self.peer_id:
            sender_index = self.sender_peer_ids.index(peer_id)
            for part_index, tensor_part in enumerate(self.parts_for_local_averaging):
                averaged_part = await self.tensor_part_reducer.accumulate_part(sender_index, part_index, tensor_part)
                self.tensor_part_container.register_processed_part(peer_index, part_index, averaged_part - tensor_part)

        else:
            code = None
            stream = self._get_peer_stub(peer_id).rpc_aggregate_part(self._generate_input_for_peer(peer_index))
            async for part_index, (averaged_part_delta, msg) in aenumerate(
                amap_in_executor(
                    lambda msg: (deserialize_torch_tensor(msg.tensor_part), msg),
                    stream,
                    max_prefetch=self.tensor_part_container.prefetch,
                )
            ):
                if code is None:
                    code = msg.code
                self.tensor_part_container.register_processed_part(peer_index, part_index, averaged_part_delta)

            if code != averaging_pb2.AVERAGED_PART:
                raise AllreduceException(
                    f"peer {peer_id} returned {averaging_pb2.MessageCode.Name(code)} "
                    f"instead of {averaging_pb2.MessageCode.Name(averaging_pb2.AVERAGED_PART)}"
                    f", allreduce failed"
                )

    async def _generate_input_for_peer(self, peer_index: int) -> AsyncIterator[averaging_pb2.AveragingData]:
        parts_aiter = self.tensor_part_container.iterate_input_parts_for(peer_index)
        first_part = await anext(parts_aiter)
        yield averaging_pb2.AveragingData(
            code=averaging_pb2.PART_FOR_AVERAGING,
            group_id=self.group_id,
            tensor_part=first_part,
        )
        async for part in parts_aiter:
            yield averaging_pb2.AveragingData(tensor_part=part)

    async def rpc_aggregate_part(
        self, stream: AsyncIterator[averaging_pb2.AveragingData], context: P2PContext
    ) -> AsyncIterator[averaging_pb2.AveragingData]:
        """a peer sends us a part of his tensor; we should average it with other peers and return the difference"""
        request: averaging_pb2.AveragingData = await anext(stream)
        reason_to_reject = self._check_reasons_to_reject(request)
        if reason_to_reject:
            yield reason_to_reject
            return

        elif request.code == averaging_pb2.PART_FOR_AVERAGING:
            try:
                sender_index = self.sender_peer_ids.index(context.remote_id)
                async for msg in self._accumulate_parts_streaming(achain(as_aiter(request), stream), sender_index):
                    yield msg

            except Exception as e:
                self.finalize(exception=e)
                yield averaging_pb2.AveragingData(code=averaging_pb2.INTERNAL_ERROR)
        else:
            error_code = averaging_pb2.MessageCode.Name(request.code)
            logger.debug(f"{self} - peer {context.remote_id} sent {error_code}, allreduce cannot continue")
            self.finalize(exception=AllreduceException(f"peer {context.remote_id} sent {error_code}."))
            yield averaging_pb2.AveragingData(code=averaging_pb2.INTERNAL_ERROR)

    def _check_reasons_to_reject(self, request: averaging_pb2.AveragingData) -> Optional[averaging_pb2.AveragingData]:
        if request.group_id != self.group_id:
            return averaging_pb2.AveragingData(code=averaging_pb2.BAD_GROUP_ID)
        elif self._future.cancelled():
            return averaging_pb2.AveragingData(code=averaging_pb2.CANCELLED)
        elif self._future.done():
            return averaging_pb2.AveragingData(code=averaging_pb2.INTERNAL_ERROR)

    async def _accumulate_parts_streaming(self, stream: AsyncIterator[averaging_pb2.AveragingData], sender_index: int):
        loop = asyncio.get_event_loop()
        async for part_index, (tensor_part, part_compression) in aenumerate(
            amap_in_executor(
                lambda msg: (deserialize_torch_tensor(msg.tensor_part), msg.tensor_part.compression),
                stream,
                max_prefetch=self.tensor_part_container.prefetch,
            )
        ):
            averaged_part = await self.tensor_part_reducer.accumulate_part(sender_index, part_index, tensor_part)

            serialized_delta = await loop.run_in_executor(
                None, lambda: serialize_torch_tensor(averaged_part - tensor_part, part_compression)
            )
            yield averaging_pb2.AveragingData(code=averaging_pb2.AVERAGED_PART, tensor_part=serialized_delta)

    async def _send_error_to_peer(self, peer_id: PeerID, code: averaging_pb2.MessageCode):
        error = averaging_pb2.AveragingData(group_id=self.group_id, code=code)
        # Coroutines are lazy, so we take the first item to start the couroutine's execution
        await afirst(self._get_peer_stub(peer_id).rpc_aggregate_part(as_aiter(error)))

    def finalize(self, *, cancel: bool = False, exception: Optional[BaseException] = None):
        """finish or terminate AllReduceRunner, propagate any errors / cancellations to peers."""
        assert not cancel or not exception, "finalize accepts either exception or cancel, but not both"
        pending_tasks = set()
        if cancel or exception:
            # propagate error to peers
            if cancel or isinstance(exception, asyncio.CancelledError):
                code = averaging_pb2.CANCELLED
            else:
                code = averaging_pb2.INTERNAL_ERROR
            logger.debug(f"{self} - notifying peers about {averaging_pb2.MessageCode.Name(code)}")
            for peer_id, mode in zip(self.ordered_peer_ids, self.modes):
                if peer_id != self.peer_id and mode != AveragingMode.CLIENT:
                    pending_tasks.add(asyncio.create_task(self._send_error_to_peer(peer_id, code)))

        if not self._future.done():
            if cancel:
                logger.debug(f"{self} - cancelled")
                self._future.cancel()
            elif exception:
                logger.debug(f"{self} - caught {exception}")
                self._future.set_exception(exception)
            else:
                logger.debug(f"{self} - finished")
                self._future.set_result(None)
            self.tensor_part_container.finalize()
            self.tensor_part_reducer.finalize()
            return pending_tasks
        else:
            logger.debug(f"{self} - could not finish: allreduce is already finished: {self._future}")
            return pending_tasks