Forward arbitrary kwargs to remote blocks

src/petals/__init__.py

@@ -24,6 +24,9 @@
     assert (
         version.parse("4.32.0") <= version.parse(transformers.__version__) < version.parse("4.35.0")
     ), "Please install a proper transformers version: pip install transformers>=4.32.0,<4.35.0"
+    assert version.parse("1.1.10") <= version.parse(
+        hivemind.__version__
+    ), "Please install a proper hivemind version: pip install hivemind>=1.1.10"
 
 
 def _override_bfloat16_mode_default():

src/petals/client/inference_session.py

@@ -4,18 +4,16 @@
 import itertools
 import time
 import uuid
-from typing import AsyncIterator, List, Optional, Tuple
+from typing import Any, AsyncIterator, Dict, List, Optional, Sequence, Tuple
 
 import torch
-from hivemind import MSGPackSerializer, anext, deserialize_torch_tensor, get_logger, serialize_torch_tensor
+from hivemind.compression import deserialize_torch_tensor, serialize_torch_tensor
 from hivemind.moe.client.remote_expert_worker import RemoteExpertWorker
-from hivemind.p2p import P2P
 from hivemind.proto import runtime_pb2
-from hivemind.utils.tensor_descr import BatchTensorDescriptor
+from hivemind.utils import MSGPackSerializer, anext, get_logger, nested_flatten
 
-from petals.client.config import ClientConfig
 from petals.client.routing import RemoteSequenceManager, maybe_log_traceback
-from petals.data_structures import CHAIN_DELIMITER, ModuleUID, RemoteSpanInfo, RPCInfo
+from petals.data_structures import CHAIN_DELIMITER, ModuleUID, RemoteSpanInfo
 from petals.server.handler import TransformerConnectionHandler
 from petals.utils.misc import DUMMY, DUMMY_INT64, is_dummy
 from petals.utils.packaging import pack_args_kwargs
@@ -32,48 +30,48 @@ class _ServerInferenceSession:
 
     def __init__(
         self,
-        config: ClientConfig,
+        sequence_manager: RemoteSequenceManager,
         span: RemoteSpanInfo,
-        uid: ModuleUID,
-        rpc_info: RPCInfo,
+        span_uids: Sequence[ModuleUID],
         inputs_queue: asyncio.Queue,
         outputs_aiter: AsyncIterator,
-        *,
+        *block_kwargs,
         max_length: int,
-        **metadata,
     ):
-        self.config = config
-        self.span, self.uid, self.rpc_info = span, uid, rpc_info
-        self.num_blocks = uid.count(CHAIN_DELIMITER) + 1
+        self.sequence_manager = sequence_manager
+        self.span, self.span_uids = span, span_uids
+        self.num_blocks = len(span_uids)
         self._inputs_queue: asyncio.Queue[runtime_pb2.ExpertRequest] = inputs_queue
         self._outputs_stream: AsyncIterator[runtime_pb2.ExpertResponse] = outputs_aiter
         self.session_id = str(uuid.uuid4())
-        self.session_metadata = dict(max_length=max_length, **metadata)
+        self.max_length = max_length
         self.stepped = False
         self.closed = False
 
         self._position = 0
         self.history = None  # Used in case of server failures to regenerate attention caches on new servers
         self.next_session = None
 
+        self.block_kwargs = block_kwargs
+        assert len(self.block_kwargs) in (0, self.num_blocks)
+
     @classmethod
     async def create(
         cls,
-        config: ClientConfig,
-        p2p: P2P,
+        sequence_manager: RemoteSequenceManager,
         span: RemoteSpanInfo,
-        uid: ModuleUID,
-        rpc_info: RPCInfo,
-        **metadata,
+        span_uids: Sequence[ModuleUID],
+        *block_kwargs: Dict[str, Any],
+        **kwargs,
     ) -> _ServerInferenceSession:
         """Create a new session for a given remote module. This code is meant to be run inside RemoteExpertWorker"""
-        stub = TransformerConnectionHandler.get_stub(p2p, span.peer_id)
+        stub = TransformerConnectionHandler.get_stub(sequence_manager.state.p2p, span.peer_id)
         inputs_queue = asyncio.Queue()
         outputs_stream = await asyncio.wait_for(
             stub.rpc_inference(cls._read_inputs_from_queue(inputs_queue)),
-            config.connect_timeout,
+            sequence_manager.config.connect_timeout,
         )
-        return cls(config, span, uid, rpc_info, inputs_queue, outputs_stream, **metadata)
+        return cls(sequence_manager, span, span_uids, inputs_queue, outputs_stream, *block_kwargs, **kwargs)
 
     @staticmethod
     async def _read_inputs_from_queue(queue: asyncio.Queue, input_timeout: Optional[float] = None) -> AsyncIterator:
@@ -84,11 +82,16 @@ async def _read_inputs_from_queue(queue: asyncio.Queue, input_timeout: Optional[
                 break  # this message means "done sending"
 
     def step(
-        self, inputs: torch.Tensor, prompts: torch.Tensor, hypo_ids: torch.LongTensor, *, step_id: str
+        self,
+        inputs: torch.Tensor,
+        prompts: Optional[torch.Tensor] = None,
+        *,
+        hypo_ids: Optional[torch.Tensor] = None,
+        step_id: str,
     ) -> torch.Tensor:
         """
         Inference step: send a chunk of input tensors and receive a chunk of outputs
-        :prompts: optional DEEP prompts, added to a prefix of each layer's outputs,
+        :param prompts: optional DEEP prompts, added to a prefix of each layer's outputs,
           if specified, deep prompts should have shape [num_layers, batch_size, prefix_len, hid_size]
         """
         if self.closed:
@@ -106,41 +109,70 @@ def step(
 
         if not self.stepped:
             inputs = self.history  # Pass full inputs including prefix
+            block_kwargs = self.block_kwargs
         else:
             inputs = inputs[:, -n_input_tokens:]  # No need to pass prefix further
+            block_kwargs = []
 
-        # serialize inputs and put them into the queue
-        input_tensors, args_structure = pack_args_kwargs(inputs, prompts, hypo_ids)
+        if prompts is None or is_dummy(prompts):
+            prompts = DUMMY
+        else:
+            assert prompts.ndim == 4, "deep prompts should have shape [num_blocks, batch_size, prefix_len, hid_size]"
+            assert prompts.shape[0] == self.num_blocks
+            assert prompts.shape[1] in (inputs.shape[0], 1)
+            assert prompts.shape[2] <= inputs.shape[1]
+            assert prompts.shape[3] == inputs.shape[2]
 
-        request_metadata = dict(session_id=self.session_id, step_id=step_id)
-        if not self.stepped:
-            request_metadata.update(self.session_metadata)
-        elif self.config.use_server_to_server:
+        if hypo_ids is None or is_dummy(hypo_ids):
+            hypo_ids = DUMMY_INT64
+        else:
+            assert len(hypo_ids) == len(inputs)
+            assert hypo_ids.dtype == torch.int64
+
+        metadata = dict(session_id=self.session_id, step_id=step_id, max_length=self.max_length)
+        metadata.update(
+            self.sequence_manager.get_request_metadata(
+                self.span.peer_id,
+                "rpc_inference",
+                self.span_uids,
+                inputs,
+                prompts,
+                *block_kwargs,
+                max_length=self.max_length,
+                session_id=self.session_id,
+                step_id=step_id,
+            )
+        )
+        if self.stepped and self.sequence_manager.config.use_server_to_server:
             next_servers = self._collect_next_servers()
             if next_servers:
-                request_metadata["next_servers"] = next_servers
+                metadata["next_servers"] = next_servers
 
-        request_metadata["args_structure"] = args_structure
+        codecs = self.sequence_manager.get_compression_codecs(
+            self.span.peer_id, "rpc_inference", self.span_uids, inputs, prompts, *block_kwargs
+        )
 
-        # TODO: make possible to use different compression method for different tensors
-        server_side_inference_schema, kwargs_schema = self.rpc_info["inference_schema"]
-        compression = server_side_inference_schema[0].compression
-        inference_schema = tuple(BatchTensorDescriptor.from_tensor(arg, compression) for arg in input_tensors)
+        # serialize inputs and put them into the queue
+        input_tensors, args_structure = pack_args_kwargs(inputs, prompts, hypo_ids, *block_kwargs)
+        args_structure = metadata.setdefault("args_structure", args_structure)
 
-        # TODO: create more explicit way to check servers schema and client's structure
-        assert len(input_tensors) >= len(
-            server_side_inference_schema
-        ), "Hidden_state, prompts and hypo_ids tensors are necessary for an inference step"
+        if codecs is None:
+            codecs = [runtime_pb2.CompressionType.NONE] * len(input_tensors)
+        else:
+            codecs = list(nested_flatten(codecs))
+            assert len(codecs) == len(
+                input_tensors
+            ), f"got {len(input_tensors)} tensors but {len(codecs)} compression codecs"
 
         outputs_serialized = RemoteExpertWorker.run_coroutine(
             self._step(
                 runtime_pb2.ExpertRequest(
-                    uid=self.uid,
+                    uid=CHAIN_DELIMITER.join(self.span_uids),
                     tensors=[
-                        serialize_torch_tensor(tensor.to(proto.dtype), proto.compression)
-                        for tensor, proto in zip(input_tensors, inference_schema)
+                        serialize_torch_tensor(tensor, compression)
+                        for tensor, compression in zip(input_tensors, codecs)
                     ],
-                    metadata=MSGPackSerializer.dumps(request_metadata),
+                    metadata=MSGPackSerializer.dumps(metadata),
                 )
             )
         )
@@ -167,7 +199,7 @@ async def _step(self, inputs_serialized: runtime_pb2.ExpertRequest) -> runtime_p
         """Inference step on serialized data. This code is meant to be run inside RemoteExpertWorker"""
         await self._inputs_queue.put(inputs_serialized)
         self.stepped = True
-        return await asyncio.wait_for(anext(self._outputs_stream), self.config.request_timeout)
+        return await asyncio.wait_for(anext(self._outputs_stream), self.sequence_manager.config.request_timeout)
 
     def close(self):
         """Finish a given inference session, close the underlying connection"""
@@ -204,14 +236,20 @@ class InferenceSession:
     An interface to a multi-step *inference* session for a sequence of remote transformer blocks
     """
 
-    def __init__(self, sequence_manager: RemoteSequenceManager, max_length: int):
+    def __init__(self, sequence_manager: RemoteSequenceManager, max_length: int, *block_kwargs: Dict[str, Any]):
         self._sequence_manager = sequence_manager
         self._closed = False
         self._server_sessions = []
         self._position = 0
         self._max_length = max_length
         self.output_ids = None
 
+        num_blocks = len(self._sequence_manager)
+        if len(block_kwargs) == 1:
+            block_kwargs = block_kwargs * num_blocks
+        assert len(block_kwargs) in (0, num_blocks), f"expected {num_blocks} block_kwargs, got {len(block_kwargs)}"
+        self.block_kwargs = block_kwargs
+
     @property
     def num_blocks(self) -> int:
         return len(self._sequence_manager)
@@ -224,17 +262,13 @@ def _enter_server_sessions(self, chosen_spans: List[RemoteSpanInfo]) -> List[_Se
         server_sessions = []
         try:
             for span in chosen_spans:
-                span_uids = CHAIN_DELIMITER.join(self._sequence_manager.block_uids[span.start : span.end])
-                metadata = self._sequence_manager.get_request_metadata("rpc_inference", span_uids, peer_id=span.peer_id)
                 session = RemoteExpertWorker.run_coroutine(
                     _ServerInferenceSession.create(
-                        self._sequence_manager.config,
-                        self._sequence_manager.state.p2p,
+                        self._sequence_manager,
                         span,
-                        span_uids,
-                        rpc_info=self._sequence_manager.rpc_info,
+                        self._sequence_manager.block_uids[span.start : span.end],
+                        *self.block_kwargs[span.start : span.end],
                         max_length=self._max_length,
-                        **metadata,
                     )
                 )
                 server_sessions.append(session)
@@ -256,8 +290,12 @@ def __enter__(self) -> "InferenceSession":
         return self
 
     def step(
-        self, inputs: torch.Tensor, prompts: Optional[torch.Tensor] = None, hypo_ids: Optional[torch.Tensor] = None
+        self,
+        inputs: torch.Tensor,
+        prompts: Optional[torch.Tensor] = None,
+        hypo_ids: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
+
         assert not self._closed
         if torch.is_grad_enabled():
             logger.warning("Running inference session with grad enabled. Gradients will *not* be propagated correctly.")
@@ -302,7 +340,10 @@ def step(
 
                     server_session = self._server_sessions[server_idx]
                     inputs = server_session.step(
-                        inputs, prompts[server_session.span.start : server_session.span.end], hypo_ids, step_id=step_id
+                        inputs,
+                        prompts[server_session.span.start : server_session.span.end],
+                        hypo_ids=hypo_ids,
+                        step_id=step_id,
                     )
 
                     server_idx += 1
@@ -328,7 +369,7 @@ def step(
         outputs = outputs.to(device=inputs_device, dtype=inputs_dtype)
         return outputs
 
-    def _update_sequence(self, server_idx: int, block_idx: int, attempt_no: int) -> int:
+    def _update_sequence(self, server_idx: int, block_idx: int, attempt_no: int):
         # If there is a failed server session, this code closes it
         self._exit_server_sessions(self._server_sessions[server_idx : server_idx + 1])