Merge branch 'main' of https://github.com/bigscience-workshop/petals into danO
Danny Boy
10 months ago
commit
15a715c5ad
@ -1,4 +1,4 @@
|
||||
from petals.client.config import ClientConfig
|
||||
from petals.client.inference_session import InferenceSession
|
||||
from petals.client.remote_sequential import RemoteSequential
|
||||
from petals.client.routing.sequence_manager import RemoteSequenceManager
|
||||
from petals.client.routing.spending_policy import NoSpendingPolicy, SpendingPolicyBase
|
||||
from petals.client.routing import NoSpendingPolicy, RemoteSequenceManager, SpendingPolicyBase
|
||||
|
||||
@ -0,0 +1,31 @@
|
||||
import dataclasses
|
||||
from typing import Optional, Sequence, Union
|
||||
|
||||
from hivemind import PeerID
|
||||
|
||||
from petals.constants import PUBLIC_INITIAL_PEERS
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class ClientConfig:
|
||||
initial_peers: Sequence[str] = tuple(PUBLIC_INITIAL_PEERS) # a list of initial peers for hivemind DHT
|
||||
dht_prefix: Optional[str] = None # a prefix for all dht keys that correspond to this model (default: model name)
|
||||
daemon_startup_timeout: int = 60 # timeout for the libp2p daemon connecting to initial peers
|
||||
|
||||
show_route: Union[str, bool] = "inference" # show chosen route through servers. one of [False, "inference", True]
|
||||
allowed_servers: Optional[Sequence[Union[PeerID, str]]] = None # if defined, send requests only to these servers
|
||||
blocked_servers: Optional[Sequence[Union[PeerID, str]]] = None # if defined, do not use these servers
|
||||
use_server_to_server: bool = True # Use direct server-to-server communication
|
||||
|
||||
connect_timeout: float = 5 # timeout for opening a connection
|
||||
request_timeout: float = 3 * 60 # timeout for forward/backward/inference requests
|
||||
update_period: float = 60 # refresh DHT information once in this many seconds
|
||||
|
||||
max_retries: Optional[int] = None # max number retries before the client raises an exception (default: inf)
|
||||
min_backoff: float = 1 # after a repeated failure, sleep for this many seconds times 2 ** (num_failures - 1)
|
||||
max_backoff: float = 60 # limit maximal sleep time between retries to this value
|
||||
ban_timeout: float = 15 # when a remote peer fails to respond, prevent routing to that peer for this many seconds
|
||||
active_adapter: Optional[str] = None # name of active LoRA adapter (usually, Hugging Face repo)
|
||||
|
||||
max_pinged: int = 3 # max servers to ping from each sequence side, per update
|
||||
ping_timeout: float = 2 # max time to wait for pings, per update
|
||||
@ -1 +1,2 @@
|
||||
"""Client-side functions responsible for choosing the best server, """
|
||||
from petals.client.routing.sequence_manager import RemoteSequenceManager, maybe_log_traceback
|
||||
from petals.client.routing.spending_policy import NoSpendingPolicy, SpendingPolicyBase
|
||||
|
||||
@ -1,124 +1,9 @@
|
||||
"""
|
||||
Utilities for declaring and retrieving active model layers using a shared DHT.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
import warnings
|
||||
|
||||
import math
|
||||
from functools import partial
|
||||
from typing import Dict, List, Optional, Sequence, Union
|
||||
warnings.warn(
|
||||
"petals.dht_utils has been moved to petals.utils.dht. This alias will be removed in Petals 2.2.0+",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
from hivemind.dht import DHT, DHTNode, DHTValue
|
||||
from hivemind.p2p import PeerID
|
||||
from hivemind.utils import DHTExpiration, MPFuture, get_dht_time, get_logger
|
||||
|
||||
from petals.data_structures import CHAIN_DELIMITER, UID_DELIMITER, ModuleUID, RemoteModuleInfo, ServerInfo
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
def declare_active_modules(
|
||||
dht: DHT,
|
||||
uids: Sequence[ModuleUID],
|
||||
server_info: ServerInfo,
|
||||
expiration_time: DHTExpiration,
|
||||
wait: bool = True,
|
||||
) -> Union[Dict[ModuleUID, bool], MPFuture[Dict[ModuleUID, bool]]]:
|
||||
"""
|
||||
Declare that your node serves the specified modules; update timestamps if declared previously
|
||||
|
||||
:param uids: a list of module ids to declare
|
||||
:param wait: if True, awaits for declaration to finish, otherwise runs in background
|
||||
:param throughput: specify your performance in terms of compute throughput
|
||||
:param expiration_time: declared modules will be visible for this many seconds
|
||||
:returns: if wait, returns store status for every key (True = store succeeded, False = store rejected)
|
||||
"""
|
||||
if isinstance(uids, str):
|
||||
uids = [uids]
|
||||
if not isinstance(uids, list):
|
||||
uids = list(uids)
|
||||
for uid in uids:
|
||||
assert isinstance(uid, ModuleUID) and UID_DELIMITER in uid and CHAIN_DELIMITER not in uid
|
||||
|
||||
return dht.run_coroutine(
|
||||
partial(_declare_active_modules, uids=uids, server_info=server_info, expiration_time=expiration_time),
|
||||
return_future=not wait,
|
||||
)
|
||||
|
||||
|
||||
async def _declare_active_modules(
|
||||
dht: DHT,
|
||||
node: DHTNode,
|
||||
uids: List[ModuleUID],
|
||||
server_info: ServerInfo,
|
||||
expiration_time: DHTExpiration,
|
||||
) -> Dict[ModuleUID, bool]:
|
||||
num_workers = len(uids) if dht.num_workers is None else min(len(uids), dht.num_workers)
|
||||
return await node.store_many(
|
||||
keys=uids,
|
||||
subkeys=[dht.peer_id.to_base58()] * len(uids),
|
||||
values=[server_info.to_tuple()] * len(uids),
|
||||
expiration_time=expiration_time,
|
||||
num_workers=num_workers,
|
||||
)
|
||||
|
||||
|
||||
def get_remote_module_infos(
|
||||
dht: DHT,
|
||||
uids: Sequence[ModuleUID],
|
||||
expiration_time: Optional[DHTExpiration] = None,
|
||||
active_adapter: Optional[str] = None,
|
||||
*,
|
||||
latest: bool = False,
|
||||
return_future: bool = False,
|
||||
) -> Union[List[Optional[RemoteModuleInfo]], MPFuture]:
|
||||
return dht.run_coroutine(
|
||||
partial(
|
||||
_get_remote_module_infos,
|
||||
uids=uids,
|
||||
active_adapter=active_adapter,
|
||||
expiration_time=expiration_time,
|
||||
latest=latest,
|
||||
),
|
||||
return_future=return_future,
|
||||
)
|
||||
|
||||
|
||||
async def _get_remote_module_infos(
|
||||
dht: DHT,
|
||||
node: DHTNode,
|
||||
uids: List[ModuleUID],
|
||||
active_adapter: Optional[str],
|
||||
expiration_time: Optional[DHTExpiration],
|
||||
latest: bool,
|
||||
) -> List[Optional[RemoteModuleInfo]]:
|
||||
if latest:
|
||||
assert expiration_time is None, "You should define either `expiration_time` or `latest`, not both"
|
||||
expiration_time = math.inf
|
||||
elif expiration_time is None:
|
||||
expiration_time = get_dht_time()
|
||||
num_workers = len(uids) if dht.num_workers is None else min(len(uids), dht.num_workers)
|
||||
found: Dict[ModuleUID, DHTValue] = await node.get_many(uids, expiration_time, num_workers=num_workers)
|
||||
|
||||
modules: List[Optional[RemoteModuleInfo]] = [None] * len(uids)
|
||||
for i, uid in enumerate(uids):
|
||||
metadata = found[uid]
|
||||
if metadata is None or not isinstance(metadata.value, dict):
|
||||
if metadata is not None:
|
||||
logger.warning(f"Incorrect metadata for {uid}: {metadata}")
|
||||
continue
|
||||
servers = {}
|
||||
for peer_id, server_info in metadata.value.items():
|
||||
try:
|
||||
peer_id = PeerID.from_base58(peer_id)
|
||||
server_info = ServerInfo.from_tuple(server_info.value)
|
||||
|
||||
if active_adapter and active_adapter not in server_info.adapters:
|
||||
logger.debug(f"Skipped server {peer_id} since it does not have adapter {active_adapter}")
|
||||
continue
|
||||
|
||||
servers[peer_id] = server_info
|
||||
except (TypeError, ValueError) as e:
|
||||
logger.warning(f"Incorrect peer entry for uid={uid}, peer_id={peer_id}: {e}")
|
||||
if servers:
|
||||
modules[i] = RemoteModuleInfo(uid, servers)
|
||||
return modules
|
||||
from petals.utils.dht import *
|
||||
|
||||
@ -0,0 +1,230 @@
|
||||
"""
|
||||
This module implements server-side computations on served blocks: forward, backward and inference; used by handler
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, AsyncIterator, Dict, Optional, Sequence, Tuple, Union
|
||||
|
||||
import torch
|
||||
from hivemind.compression.serialization import deserialize_torch_tensor, serialize_torch_tensor
|
||||
from hivemind.moe.expert_uid import ExpertUID
|
||||
from hivemind.proto import runtime_pb2
|
||||
from hivemind.utils.logging import get_logger
|
||||
from hivemind.utils.nested import nested_flatten
|
||||
|
||||
from petals.data_structures import Handle, InferenceMetadata
|
||||
from petals.server.backend import TransformerBackend
|
||||
from petals.server.task_pool import PrioritizedTaskPool
|
||||
from petals.server.task_prioritizer import TaskPrioritizerBase
|
||||
from petals.utils.convert_block import QuantType
|
||||
from petals.utils.misc import DUMMY, is_dummy
|
||||
from petals.utils.packaging import unpack_args_kwargs
|
||||
|
||||
# We prioritize short inference requests and make them use a *merged* inference pool,
|
||||
# so they are processed without interruptions and extra overheads
|
||||
# TODO: Increase the NF4 threshold once bitsandbytes ships efficient NF4 kernel for parallel forward
|
||||
MAX_SHORT_INFERENCE_TOKENS = 128
|
||||
MAX_NF4_SHORT_INFERENCE_TOKENS = 1
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
async def run_rpc_forward(
|
||||
*flat_tensors: torch.Tensor,
|
||||
requested_backends: Sequence[TransformerBackend],
|
||||
active_adapter: str = "",
|
||||
prioritizer: TaskPrioritizerBase,
|
||||
points: int = 0,
|
||||
args_structure: Any = None,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Run forward pass on deserialized inputs and prompts, used by rpc_forward and rpc_forward_stream
|
||||
|
||||
:param flat_tensors: a list of tensors that includes first layer inputs, optional prompts and extra tensors
|
||||
:note: some input tensors can be missing, in which case they will be replaced with dummy tensors (see is_dummy)
|
||||
:param requested_backends: a sequence of transformer blocks in the same order as they appear in forward pass
|
||||
:returns: hidden states after the last layer [batch_size, seq_length, hid_size]
|
||||
"""
|
||||
if args_structure is not None:
|
||||
# TODO: kwargs currently is unused, it can be used later for peft-like adaptation
|
||||
flat_tensors, kwargs = unpack_args_kwargs(flat_tensors, args_structure)
|
||||
hidden_states, prompts, *_ = flat_tensors
|
||||
|
||||
dtype = requested_backends[0].dtype
|
||||
# check parse input tensors and cast dtypes
|
||||
hidden_states = hidden_states.to(dtype)
|
||||
assert hidden_states.ndim == 3
|
||||
if prompts is None or is_dummy(prompts):
|
||||
prompts = [DUMMY] * len(requested_backends)
|
||||
else:
|
||||
prompts = [p.squeeze(0) for p in prompts.to(requested_backends[0].dtype).split(1, dim=0)]
|
||||
|
||||
# Run a chain of requested backends
|
||||
for backend, prompt in zip(requested_backends, prompts):
|
||||
if not is_dummy(prompt):
|
||||
hidden_states[:, : prompt.shape[1]] += prompt
|
||||
|
||||
assert isinstance(backend.inference_pool, PrioritizedTaskPool), "petals support only prioritized pools"
|
||||
priority = prioritizer.prioritize(
|
||||
hidden_states, points=points / len(requested_backends), backend=backend, type="forward"
|
||||
)
|
||||
(hidden_states,) = await backend.forward_pool.submit_task(
|
||||
hidden_states,
|
||||
active_adapter,
|
||||
priority=priority,
|
||||
)
|
||||
assert isinstance(hidden_states, torch.Tensor)
|
||||
assert (
|
||||
hidden_states.ndim == 3
|
||||
), f"inputs to {type(backend)} must be a list with a single 3d tensor of hidden states"
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
async def run_rpc_backward(
|
||||
*flat_tensors: torch.Tensor,
|
||||
requested_backends: Sequence[TransformerBackend],
|
||||
active_adapter: str = "",
|
||||
prioritizer: TaskPrioritizerBase,
|
||||
points: int = 0,
|
||||
args_structure: Any = None,
|
||||
) -> Union[torch.Tensor, Sequence[torch.Tensor]]:
|
||||
if args_structure is not None:
|
||||
# TODO: kwargs currently is unused, it can be used later for peft-like adaptation
|
||||
flat_tensors, kwargs = unpack_args_kwargs(flat_tensors, args_structure)
|
||||
inputs, grad_outputs, prompts, *_ = flat_tensors
|
||||
|
||||
# Cast inputs & grad outputs to backend dtype
|
||||
inputs = inputs.to(requested_backends[0].dtype)
|
||||
grad_outputs = grad_outputs.to(requested_backends[-1].dtype)
|
||||
|
||||
if prompts is None or is_dummy(prompts):
|
||||
prompts = [DUMMY] * len(requested_backends)
|
||||
else:
|
||||
prompts = [p.squeeze(0) for p in prompts.to(requested_backends[0].dtype).split(1, dim=0)]
|
||||
|
||||
# Run a forward chain to collect intermediate inputs
|
||||
# Note that we do not forward for the last module since we do not need its output
|
||||
inter_inputs = []
|
||||
for backend, prompt in zip(requested_backends[:-1], prompts[:-1]):
|
||||
assert inputs.ndim == 3, f"inputs to {type(backend)} must be a single 3d tensor of hidden states"
|
||||
if not is_dummy(prompt):
|
||||
inputs[:, : prompt.shape[1]] += prompt
|
||||
inter_inputs.append(inputs)
|
||||
assert isinstance(backend.inference_pool, PrioritizedTaskPool), "petals support only prioritized pools"
|
||||
priority = prioritizer.prioritize(
|
||||
inputs, points=points / len(requested_backends), backend=backend, type="forward_in_backward"
|
||||
)
|
||||
(inputs,) = await backend.forward_pool.submit_task(inputs, active_adapter, priority=priority)
|
||||
|
||||
assert isinstance(inputs, torch.Tensor)
|
||||
|
||||
if not is_dummy(prompts[-1]):
|
||||
inputs[:, : prompts[-1].shape[1]] += prompts[-1]
|
||||
inter_inputs.append(inputs)
|
||||
|
||||
assert len(inter_inputs) == len(prompts) == len(requested_backends), "internal shape error during backward"
|
||||
grad_prompts_reversed = []
|
||||
# Run a chain of requested backends
|
||||
for inp, prompt, backend in zip(*map(reversed, (inter_inputs, prompts, requested_backends))):
|
||||
assert isinstance(backend.inference_pool, PrioritizedTaskPool), "petals support only prioritized pools"
|
||||
priority = prioritizer.prioritize(
|
||||
inp, grad_outputs, points=points / len(requested_backends), backend=backend, type="backward"
|
||||
)
|
||||
(grad_outputs,) = await backend.backward_pool.submit_task(inp, grad_outputs, active_adapter, priority=priority)
|
||||
|
||||
assert isinstance(grad_outputs, torch.Tensor)
|
||||
if not is_dummy(prompt):
|
||||
grad_prompts_reversed.append(grad_outputs[:, : prompt.shape[1]].unsqueeze(0))
|
||||
|
||||
grad_prompts = torch.cat(grad_prompts_reversed[::-1], dim=0) if grad_prompts_reversed else DUMMY
|
||||
return [grad_outputs] if is_dummy(grad_prompts) else [grad_outputs, grad_prompts] # TODO un-duct-tape
|
||||
|
||||
|
||||
async def iterate_rpc_inference(
|
||||
requested_uids: Sequence[ExpertUID],
|
||||
requested_backends: Sequence[TransformerBackend],
|
||||
active_adapter: Optional[str],
|
||||
input_iterator: AsyncIterator[Tuple[runtime_pb2.ExpertRequest, dict]],
|
||||
cache_handles: Sequence[Sequence[Handle]],
|
||||
*,
|
||||
max_length: int,
|
||||
prioritizer: TaskPrioritizerBase,
|
||||
points: int,
|
||||
quant_type: QuantType,
|
||||
args_structure: Any = None,
|
||||
) -> AsyncIterator[Tuple[Sequence[runtime_pb2.Tensor], bool]]:
|
||||
assert len(cache_handles) == len(requested_backends)
|
||||
|
||||
prefix_length = 0
|
||||
point_per_piece = points / max_length if max_length > 0 else 0.0
|
||||
|
||||
async for request, step_metadata in input_iterator:
|
||||
flat_tensors = tuple(deserialize_torch_tensor(tensor) for tensor in request.tensors)
|
||||
if args_structure is not None:
|
||||
# TODO: kwargs currently is unused, it can be used later for peft-like adaptation
|
||||
flat_tensors, kwargs = unpack_args_kwargs(flat_tensors, args_structure)
|
||||
|
||||
hidden_states, prompts, hypo_ids, *_ = flat_tensors
|
||||
batch_size, length_increment, _ = hidden_states.shape
|
||||
|
||||
# Cast inputs to backend dtype
|
||||
hidden_states = hidden_states.to(requested_backends[0].dtype)
|
||||
assert hypo_ids.dtype == torch.int64, f"hypo ids must be int64, got {hypo_ids.dtype}"
|
||||
|
||||
# parse deep prompts (optional argument)
|
||||
has_prompts = prompts is not None and not is_dummy(prompts)
|
||||
if not has_prompts:
|
||||
prompts = [None] * len(requested_backends)
|
||||
else:
|
||||
prompts = [p.squeeze(0) for p in prompts.to(requested_backends[0].dtype).split(1, dim=0)]
|
||||
prompts = [prompt if not is_dummy(prompt) else None for prompt in prompts]
|
||||
|
||||
if not (len(requested_backends) == len(prompts)):
|
||||
raise ValueError(f"Received {len(prompts)} prompts for {len(requested_backends)} backends")
|
||||
|
||||
if prefix_length + length_increment > max_length:
|
||||
raise ValueError(
|
||||
f"Maximum length exceeded: prefix {prefix_length} + current {length_increment}"
|
||||
f" exceeds pre-allocated maximum {max_length}"
|
||||
)
|
||||
|
||||
merge_max_tokens = MAX_NF4_SHORT_INFERENCE_TOKENS if quant_type == QuantType.NF4 else MAX_SHORT_INFERENCE_TOKENS
|
||||
can_merge_pools = batch_size * length_increment <= merge_max_tokens
|
||||
priority = prioritizer.prioritize(
|
||||
hidden_states,
|
||||
hypo_ids,
|
||||
points=point_per_piece,
|
||||
requested_uids=requested_uids,
|
||||
type="short_inference" if can_merge_pools else "inference",
|
||||
)
|
||||
|
||||
# A client may pass a tensor with 0 tokens. This is a special case that occurs, e.g.
|
||||
# when user wants to pre-allocate cache or check that server *can* allocate that cache.
|
||||
if hidden_states.numel() > 0:
|
||||
assert hidden_states.ndim == 3, f"hidden states must be a single 3d tensor"
|
||||
if can_merge_pools:
|
||||
inference_infos = tuple(
|
||||
InferenceMetadata(uid, prefix_length, tuple(handles), active_adapter)
|
||||
for uid, handles in zip(requested_uids, cache_handles)
|
||||
)
|
||||
(hidden_states,) = await requested_backends[0].inference_pool.submit_task(
|
||||
hidden_states, hypo_ids, inference_infos, *prompts, priority=priority
|
||||
)
|
||||
else:
|
||||
for backend, uid, handles, prompt in zip(requested_backends, requested_uids, cache_handles, prompts):
|
||||
inference_infos = (InferenceMetadata(uid, prefix_length, tuple(handles), active_adapter),)
|
||||
(hidden_states,) = await backend.inference_pool.submit_task(
|
||||
hidden_states, hypo_ids, inference_infos, prompt, priority=priority
|
||||
)
|
||||
|
||||
# serialize and send last layer outputs
|
||||
output_tensors = [
|
||||
serialize_torch_tensor(result.to(proto.dtype), proto.compression, allow_inplace=True)
|
||||
for result, proto in zip((hidden_states,), nested_flatten(requested_backends[-1].outputs_schema))
|
||||
]
|
||||
can_push = not has_prompts
|
||||
yield output_tensors, can_push
|
||||
|
||||
# prepare for next step
|
||||
prefix_length += length_increment
|
||||
@ -0,0 +1,124 @@
|
||||
"""
|
||||
Utilities for declaring and retrieving active model layers using a shared DHT.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from functools import partial
|
||||
from typing import Dict, List, Optional, Sequence, Union
|
||||
|
||||
from hivemind.dht import DHT, DHTNode, DHTValue
|
||||
from hivemind.p2p import PeerID
|
||||
from hivemind.utils import DHTExpiration, MPFuture, get_dht_time, get_logger
|
||||
|
||||
from petals.data_structures import CHAIN_DELIMITER, UID_DELIMITER, ModuleUID, RemoteModuleInfo, ServerInfo
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
def declare_active_modules(
|
||||
dht: DHT,
|
||||
uids: Sequence[ModuleUID],
|
||||
server_info: ServerInfo,
|
||||
expiration_time: DHTExpiration,
|
||||
wait: bool = True,
|
||||
) -> Union[Dict[ModuleUID, bool], MPFuture[Dict[ModuleUID, bool]]]:
|
||||
"""
|
||||
Declare that your node serves the specified modules; update timestamps if declared previously
|
||||
|
||||
:param uids: a list of module ids to declare
|
||||
:param wait: if True, awaits for declaration to finish, otherwise runs in background
|
||||
:param throughput: specify your performance in terms of compute throughput
|
||||
:param expiration_time: declared modules will be visible for this many seconds
|
||||
:returns: if wait, returns store status for every key (True = store succeeded, False = store rejected)
|
||||
"""
|
||||
if isinstance(uids, str):
|
||||
uids = [uids]
|
||||
if not isinstance(uids, list):
|
||||
uids = list(uids)
|
||||
for uid in uids:
|
||||
assert isinstance(uid, ModuleUID) and UID_DELIMITER in uid and CHAIN_DELIMITER not in uid
|
||||
|
||||
return dht.run_coroutine(
|
||||
partial(_declare_active_modules, uids=uids, server_info=server_info, expiration_time=expiration_time),
|
||||
return_future=not wait,
|
||||
)
|
||||
|
||||
|
||||
async def _declare_active_modules(
|
||||
dht: DHT,
|
||||
node: DHTNode,
|
||||
uids: List[ModuleUID],
|
||||
server_info: ServerInfo,
|
||||
expiration_time: DHTExpiration,
|
||||
) -> Dict[ModuleUID, bool]:
|
||||
num_workers = len(uids) if dht.num_workers is None else min(len(uids), dht.num_workers)
|
||||
return await node.store_many(
|
||||
keys=uids,
|
||||
subkeys=[dht.peer_id.to_base58()] * len(uids),
|
||||
values=[server_info.to_tuple()] * len(uids),
|
||||
expiration_time=expiration_time,
|
||||
num_workers=num_workers,
|
||||
)
|
||||
|
||||
|
||||
def get_remote_module_infos(
|
||||
dht: DHT,
|
||||
uids: Sequence[ModuleUID],
|
||||
expiration_time: Optional[DHTExpiration] = None,
|
||||
active_adapter: Optional[str] = None,
|
||||
*,
|
||||
latest: bool = False,
|
||||
return_future: bool = False,
|
||||
) -> Union[List[Optional[RemoteModuleInfo]], MPFuture]:
|
||||
return dht.run_coroutine(
|
||||
partial(
|
||||
_get_remote_module_infos,
|
||||
uids=uids,
|
||||
active_adapter=active_adapter,
|
||||
expiration_time=expiration_time,
|
||||
latest=latest,
|
||||
),
|
||||
return_future=return_future,
|
||||
)
|
||||
|
||||
|
||||
async def _get_remote_module_infos(
|
||||
dht: DHT,
|
||||
node: DHTNode,
|
||||
uids: List[ModuleUID],
|
||||
active_adapter: Optional[str],
|
||||
expiration_time: Optional[DHTExpiration],
|
||||
latest: bool,
|
||||
) -> List[Optional[RemoteModuleInfo]]:
|
||||
if latest:
|
||||
assert expiration_time is None, "You should define either `expiration_time` or `latest`, not both"
|
||||
expiration_time = math.inf
|
||||
elif expiration_time is None:
|
||||
expiration_time = get_dht_time()
|
||||
num_workers = len(uids) if dht.num_workers is None else min(len(uids), dht.num_workers)
|
||||
found: Dict[ModuleUID, DHTValue] = await node.get_many(uids, expiration_time, num_workers=num_workers)
|
||||
|
||||
modules: List[Optional[RemoteModuleInfo]] = [None] * len(uids)
|
||||
for i, uid in enumerate(uids):
|
||||
metadata = found[uid]
|
||||
if metadata is None or not isinstance(metadata.value, dict):
|
||||
if metadata is not None:
|
||||
logger.warning(f"Incorrect metadata for {uid}: {metadata}")
|
||||
continue
|
||||
servers = {}
|
||||
for peer_id, server_info in metadata.value.items():
|
||||
try:
|
||||
peer_id = PeerID.from_base58(peer_id)
|
||||
server_info = ServerInfo.from_tuple(server_info.value)
|
||||
|
||||
if active_adapter and active_adapter not in server_info.adapters:
|
||||
logger.debug(f"Skipped server {peer_id} since it does not have adapter {active_adapter}")
|
||||
continue
|
||||
|
||||
servers[peer_id] = server_info
|
||||
except (TypeError, ValueError) as e:
|
||||
logger.warning(f"Incorrect peer entry for uid={uid}, peer_id={peer_id}: {e}")
|
||||
if servers:
|
||||
modules[i] = RemoteModuleInfo(uid, servers)
|
||||
return modules
|
||||
@ -0,0 +1,49 @@
|
||||
from typing import Any, Dict, List, Tuple
|
||||
|
||||
import torch
|
||||
from hivemind import nested_flatten, nested_pack
|
||||
|
||||
# TODO: Move functions to hivemind
|
||||
|
||||
|
||||
def _mark_masked_tensor(index: int) -> bytes:
|
||||
return b"__T" + str(index).encode()
|
||||
|
||||
|
||||
def _is_masked_tensor(item: Any) -> bool:
|
||||
return isinstance(item, bytes) and item.startswith(b"__T")
|
||||
|
||||
|
||||
def _get_tensor_index(item: bytes) -> int:
|
||||
return int(item[3:])
|
||||
|
||||
|
||||
def pack_args_kwargs(*args, **kwargs) -> Tuple[List[torch.Tensor], Any]:
|
||||
"""
|
||||
Check the function's arguments and pack all tensors into different flattened lists.
|
||||
:returns: a flattened list of tensors and args and kwargs, where tensors were masked
|
||||
"""
|
||||
masked_flat_values, flat_tensors, tensor_to_index = [], [], {}
|
||||
for value in nested_flatten((args, kwargs)):
|
||||
if isinstance(value, torch.Tensor):
|
||||
tensor_index = tensor_to_index.setdefault(value, len(flat_tensors))
|
||||
if tensor_index == len(flat_tensors):
|
||||
flat_tensors.append(value)
|
||||
masked_flat_values.append(_mark_masked_tensor(tensor_index))
|
||||
else:
|
||||
masked_flat_values.append(value)
|
||||
return flat_tensors, nested_pack(masked_flat_values, (args, kwargs))
|
||||
|
||||
|
||||
def unpack_args_kwargs(flat_tensors: List[torch.Tensor], args_structure: Any):
|
||||
"""
|
||||
Restore arguments after `pack_args_kwargs` function.
|
||||
:returns: list of args and dict of kwargs
|
||||
"""
|
||||
return nested_pack(
|
||||
(
|
||||
value if not _is_masked_tensor(value) else flat_tensors[_get_tensor_index(value)]
|
||||
for value in nested_flatten(args_structure)
|
||||
),
|
||||
args_structure,
|
||||
)
|
||||
Binary file not shown.
Loading…
Reference in New Issue