Merge branch 'main' of https://github.com/artek0chumak/petals
artek0chumak
9 months ago
commit
d1f0a320c7
@ -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,349 +1,142 @@
|
||||
import contextlib
|
||||
from typing import List, Optional
|
||||
import dataclasses
|
||||
from contextvars import ContextVar
|
||||
from typing import ContextManager, List, Optional
|
||||
|
||||
import torch
|
||||
import transformers
|
||||
from hivemind.utils.logging import get_logger
|
||||
from transformers.generation.utils import ModelOutput
|
||||
|
||||
from petals.client.inference_session import InferenceSession
|
||||
from petals.utils.generation_algorithms import (
|
||||
BeamSearchAlgorithm,
|
||||
DecodingAlgorithm,
|
||||
GreedyAlgorithm,
|
||||
NucleusAlgorithm,
|
||||
SamplingAlgorithm,
|
||||
TopKAlgorithm,
|
||||
)
|
||||
from petals.utils.generation_constraints import ABCBloomConstraint, EosConstraint
|
||||
from petals.client.remote_sequential import RemoteSequential
|
||||
from petals.utils.misc import DUMMY, docstring_from
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
class RemoteGenerationMixin:
|
||||
"""
|
||||
A class containing all functions for auto-regressive text generation, to be used as a mixin in [`BloomForCausalLM`].
|
||||
The class exposes can be used for:
|
||||
- *greedy decoding*.
|
||||
- *multinomial, top-k and top-p sampling*.
|
||||
- *beam-search decoding*
|
||||
|
||||
This class is similar to transformer's [`generation_utils.GenerationMixin`], it can be used instead of it.
|
||||
However, it has some differences for remote usage.
|
||||
"""
|
||||
|
||||
def inference_session(self, **kwargs) -> InferenceSession:
|
||||
"""
|
||||
Returns an inference session for the model's RemoteSequential module.
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class RemotePastKeyValues:
|
||||
"""A mock class representing the fact that `past_key_values` do exist but are stored on remote servers."""
|
||||
|
||||
:param max_length: Maximal expected length of inference results. Servers use this parameter
|
||||
to calculate the size of attention caches allocated to this client.
|
||||
"""
|
||||
hypo_ids: Optional[torch.LongTensor] = None
|
||||
|
||||
return self.transformer.h.inference_session(**kwargs)
|
||||
def __getitem__(self, _index: int) -> List[torch.Tensor]:
|
||||
return [DUMMY] # For compatibility with BloomForCausalLM.prepare_inputs_for_generation()
|
||||
|
||||
@torch.inference_mode()
|
||||
def generate(
|
||||
self,
|
||||
inputs: Optional[torch.Tensor] = None,
|
||||
*,
|
||||
do_sample: Optional[bool] = None,
|
||||
temperature: float = 1.0,
|
||||
top_k: Optional[int] = None,
|
||||
top_p: Optional[float] = None,
|
||||
num_beams: Optional[int] = 1,
|
||||
bos_token_id: Optional[int] = None,
|
||||
eos_token_id: Optional[int] = None,
|
||||
pad_token_id: Optional[int] = None,
|
||||
max_length: Optional[int] = None,
|
||||
max_new_tokens: Optional[int] = None,
|
||||
decoding_algorithm: Optional[DecodingAlgorithm] = None,
|
||||
provided_constraints: List[ABCBloomConstraint] = [],
|
||||
num_return_sequences: Optional[int] = None,
|
||||
session: Optional[InferenceSession] = None,
|
||||
) -> torch.LongTensor:
|
||||
"""
|
||||
Generates sequences of token ids for models with a language modeling head.
|
||||
|
||||
:param inputs: The input tokens to the model.
|
||||
:param do_sample: Whether to sample from the model predictions or take the argmax.
|
||||
:param temperature: The temperature to use for sampling.
|
||||
:param top_k: The number of results to return.
|
||||
:param top_p: The cumulative probability of results to return.
|
||||
:param num_beams: The number of beams to use for beam search.
|
||||
:param bos_token_id: The id of the beginning of sentence token.
|
||||
:param eos_token_id: The id of the end of sentence token.
|
||||
:param pad_token_id: The id of the padding token.
|
||||
:param max_length: The maximum number of tokens in the output (including input tokens).
|
||||
:param max_new_tokens: The maximum number of tokens to generate.
|
||||
:param decoding_algorithm: The decoding algorithm to use.
|
||||
:param provided_constraints: A list of constraints to use.
|
||||
:param num_return_sequences: How many hypothesis from the beam will be in output.
|
||||
"""
|
||||
_skipped_tokens = ContextVar("skipped_tokens", default=0)
|
||||
|
||||
prefix_length = 0 if inputs is None else inputs.size(1)
|
||||
prefix_length += self.config.pre_seq_len
|
||||
|
||||
bos_token_id = bos_token_id if bos_token_id is not None else self.config.bos_token_id
|
||||
pad_token_id = pad_token_id if pad_token_id is not None else self.config.pad_token_id
|
||||
eos_token_id = eos_token_id if eos_token_id is not None else self.config.eos_token_id
|
||||
class _SkipTokensMixin:
|
||||
# This override is used in RemoteGenerationMixin by has to be defined in a class not named as "GenerationMixin"
|
||||
# due to how transformers.PreTrainedModel.can_generate() works
|
||||
def prepare_inputs_for_generation(self, input_ids: torch.LongTensor, **kwargs) -> dict:
|
||||
input_ids = input_ids[:, _skipped_tokens.get() :]
|
||||
_skipped_tokens.set(0)
|
||||
return super().prepare_inputs_for_generation(input_ids, **kwargs)
|
||||
|
||||
assert (max_length is None) != (max_new_tokens is None), "please set max_length or max_new_tokens (not both)"
|
||||
if max_length is not None and max_new_tokens is None:
|
||||
max_new_tokens = max_length - prefix_length
|
||||
assert max_new_tokens > 0, f"Provided max_length is less than prefix size: {max_length} < {inputs.size(1)}"
|
||||
elif max_length is None and max_new_tokens is not None:
|
||||
max_length = prefix_length + max_new_tokens
|
||||
|
||||
resuming_session = session is not None and session.last_token_id is not None
|
||||
if num_beams > 1 and resuming_session:
|
||||
raise NotImplementedError(
|
||||
"Resuming inference session in .generate() along with beam search is not supported yet"
|
||||
)
|
||||
class RemoteGenerationMixin(_SkipTokensMixin):
|
||||
"""
|
||||
This class is an upgrade to `transformers.GenerationMixin` that:
|
||||
|
||||
- Designed to be compatible with most `transformers.GenerationMixin` strategies and options
|
||||
- Supports generation inside a remote InferenceSession, so that remote servers store your attention caches and
|
||||
you don't have to rerun the prefix through all the servers to generate each new token
|
||||
- Supports multiple `.generate()` calls inside one InferenceSession, so you can easily run interactive generation
|
||||
by showing tokens on the fly (multiple calls like `.generate(None, max_new_tokens=1, ...)`) or
|
||||
accept prompts from a user in a chat bot (multiple calls like `.generate(new_prompts, ...)`).
|
||||
- If there is no active session, `.generate()` will create a new InferenceSession with proper `max_length`.
|
||||
Otherwise, `.generate()` will use the active session. You can use the `session=...` argument to override that.
|
||||
"""
|
||||
|
||||
if inputs is not None:
|
||||
assert isinstance(inputs, torch.Tensor) and inputs.ndim == 2, "inputs must be a 2d tensor [batch, length]"
|
||||
if resuming_session:
|
||||
inputs = torch.cat([session.last_token_id, inputs], dim=1)
|
||||
else:
|
||||
if resuming_session:
|
||||
inputs = session.last_token_id
|
||||
else:
|
||||
assert bos_token_id is not None, "You have to provide a bos_token_id if you do not provide inputs"
|
||||
inputs = torch.tensor([[bos_token_id]] * num_beams, dtype=torch.long, device=self.device)
|
||||
batch_size = inputs.size(0)
|
||||
@docstring_from(RemoteSequential.active_session)
|
||||
@property
|
||||
def active_session(self) -> Optional[InferenceSession]:
|
||||
return self.transformer.h.active_session
|
||||
|
||||
if decoding_algorithm is None:
|
||||
if do_sample:
|
||||
decoding_algorithm = self._choose_sample_algorithm(temperature, top_k, top_p)
|
||||
elif num_beams is not None and num_beams > 1:
|
||||
decoding_algorithm = BeamSearchAlgorithm(num_beams, batch_size=batch_size)
|
||||
else:
|
||||
if top_k is not None or top_p is not None:
|
||||
logger.warning("You passed top_k or top_p but did not pass do_sample=True. Running greedy sampling")
|
||||
decoding_algorithm = GreedyAlgorithm()
|
||||
@docstring_from(RemoteSequential.use_session)
|
||||
def use_session(self, session: Optional[InferenceSession]) -> ContextManager[InferenceSession]:
|
||||
return self.transformer.h.use_session(session)
|
||||
|
||||
if num_beams > 1:
|
||||
inputs = torch.cat([inputs] * num_beams, dim=0)
|
||||
if batch_size > 1:
|
||||
# TODO: resolve padding problem
|
||||
logger.warning(
|
||||
f"You set batch_size {batch_size} within beam search generation. "
|
||||
f"Be careful, results on sequences with different length may be padded wrong way"
|
||||
)
|
||||
@docstring_from(RemoteSequential.inference_session)
|
||||
def inference_session(self, **kwargs) -> ContextManager[InferenceSession]:
|
||||
return self.transformer.h.inference_session(**kwargs)
|
||||
|
||||
if num_return_sequences is None:
|
||||
num_return_sequences = 1
|
||||
@docstring_from(transformers.GenerationMixin.generate.__doc__)
|
||||
def generate(
|
||||
self, inputs: Optional[torch.Tensor] = None, *args, session: Optional[InferenceSession] = None, **kwargs
|
||||
):
|
||||
self._fix_generate_kwargs(kwargs)
|
||||
|
||||
if session is not None:
|
||||
# If a session specified explicitly, use it
|
||||
context_manager = self.use_session(session)
|
||||
elif self.active_session is not None:
|
||||
# If there's an active session, don't do anything
|
||||
context_manager = contextlib.nullcontext(self.active_session)
|
||||
else:
|
||||
# If there's no active session, create a new one
|
||||
|
||||
assert num_return_sequences <= num_beams, (
|
||||
f"You want more sequences than the beam has."
|
||||
" Check num_return_sequences: {num_return_sequences} and num_beams: {num_beams}."
|
||||
)
|
||||
max_length = kwargs.get("max_length")
|
||||
max_new_tokens = kwargs.get("max_new_tokens")
|
||||
assert (max_length is None) != (
|
||||
max_new_tokens is None
|
||||
), "You should set `max_length` or `max_new_tokens` (but not both) to reserve server-side attention caches"
|
||||
|
||||
constraints = self._get_constraints(
|
||||
inputs=inputs,
|
||||
eos_token_id=eos_token_id,
|
||||
pad_token_id=pad_token_id,
|
||||
provided_constraints=provided_constraints,
|
||||
)
|
||||
if max_length is not None:
|
||||
session_max_length = max_length
|
||||
else:
|
||||
session_max_length = (inputs.shape[1] if inputs is not None else 0) + max_new_tokens
|
||||
context_manager = self.inference_session(max_length=session_max_length)
|
||||
|
||||
if session is None:
|
||||
context_manager = self.inference_session(max_length=max_length)
|
||||
else:
|
||||
context_manager = contextlib.nullcontext(session) # Doesn't actually enter session or exit from it
|
||||
with context_manager as session:
|
||||
outputs = []
|
||||
# Find samples with padded inputs.
|
||||
# They will be changed before all of the samples have right length.
|
||||
if torch.any(inputs == pad_token_id): # TODO: move to prepare_inputs
|
||||
outputs += [inputs[:, : inputs.size(1) - (inputs == pad_token_id).sum(-1).max()]]
|
||||
# Prepend the tokens from the previous .generate() call
|
||||
n_prev_tokens = session.output_ids.shape[1] if session.output_ids is not None else 0
|
||||
if n_prev_tokens > 0:
|
||||
if kwargs.get("num_beams", 1) > 1:
|
||||
logger.warning(
|
||||
"Beam search will not work properly in the resumed petals.InferenceSession "
|
||||
"since intermediate beam entries are lost"
|
||||
)
|
||||
|
||||
if inputs is not None:
|
||||
inputs = torch.cat([session.output_ids, inputs], dim=1)
|
||||
else:
|
||||
inputs = session.output_ids
|
||||
|
||||
# Don't actually run all previous tokens through the transformer,
|
||||
# but keep them for transformers.GenerationMixin (e.g., to compute repetition_penalty)
|
||||
_skipped_tokens.set(max(0, n_prev_tokens - 1))
|
||||
|
||||
result = super().generate(inputs, *args, **kwargs)
|
||||
|
||||
sequences = result.sequences if isinstance(result, ModelOutput) else result
|
||||
# Save tokens from this .generate() call
|
||||
session.output_ids = sequences
|
||||
# Crop the last tokens from the previous call
|
||||
sequences = sequences[:, n_prev_tokens:].clone()
|
||||
if isinstance(result, ModelOutput):
|
||||
result.sequences = sequences
|
||||
else:
|
||||
outputs += [inputs]
|
||||
last_token_id = None
|
||||
seq_idx = outputs[0].size(1)
|
||||
hypo_ids = torch.arange(outputs[0].size(0))
|
||||
while True:
|
||||
hidden_state = self.transformer.word_embeddings(outputs[-1])
|
||||
intermediate_prompts = None
|
||||
if self.config.pre_seq_len > 0 and len(outputs) == 1:
|
||||
prompts, intermediate_prompts = self.transformer.get_prompt(hidden_state.size(0))
|
||||
hidden_state = torch.cat([prompts, hidden_state], dim=1)
|
||||
hidden_state = self.transformer.word_embeddings_layernorm(hidden_state)
|
||||
result = sequences
|
||||
|
||||
hidden_state = session.step(hidden_state, prompts=intermediate_prompts, hypo_ids=hypo_ids)[:, -1]
|
||||
return result
|
||||
|
||||
hidden_state = self.transformer.ln_f(hidden_state)
|
||||
lm_logits = self.lm_head(hidden_state)
|
||||
@staticmethod
|
||||
def _fix_generate_kwargs(kwargs: dict) -> dict:
|
||||
# Suppress inappropriate "Both max_new_tokens and max_length" HF warning
|
||||
if "max_length" in kwargs and kwargs["max_length"] is None:
|
||||
del kwargs["max_length"]
|
||||
|
||||
for constraint in constraints:
|
||||
lm_logits = constraint(last_token_id, lm_logits, hypo_ids)
|
||||
last_token_id, hypo_ids = decoding_algorithm(lm_logits)
|
||||
# Support do_sample = {0, 1} for backward compatibility with Petals < 2.1.0
|
||||
do_sample = kwargs.get("do_sample")
|
||||
if isinstance(do_sample, int):
|
||||
kwargs["do_sample"] = bool(do_sample)
|
||||
|
||||
# If some samples were padded, change only these samples
|
||||
if seq_idx < inputs.size(1):
|
||||
pad_token_mask = inputs[:, seq_idx : seq_idx + 1] == pad_token_id
|
||||
last_token_id = (~pad_token_mask) * inputs[
|
||||
:, seq_idx : seq_idx + 1
|
||||
] + pad_token_mask * last_token_id
|
||||
|
||||
# TODO: refactor outputs
|
||||
if num_beams > 1:
|
||||
for i in range(len(outputs), 1, -1):
|
||||
outputs[i - 1] = outputs[i - 1][hypo_ids]
|
||||
|
||||
outputs.append(last_token_id)
|
||||
session.last_token_id = last_token_id
|
||||
seq_idx += 1
|
||||
if torch.all(last_token_id == eos_token_id) or len(outputs) > max_new_tokens:
|
||||
break
|
||||
|
||||
outputs = torch.cat(outputs, dim=-1)
|
||||
|
||||
if resuming_session:
|
||||
outputs = outputs[:, 1:]
|
||||
if num_beams > 1:
|
||||
pre_return_idx = [
|
||||
torch.arange(idx, num_return_sequences * batch_size, batch_size) for idx in range(batch_size)
|
||||
]
|
||||
return_idx = torch.cat(pre_return_idx, dim=0)
|
||||
outputs = outputs[return_idx]
|
||||
|
||||
return outputs
|
||||
|
||||
def greedy_search(
|
||||
self,
|
||||
input_ids: torch.LongTensor,
|
||||
max_length: Optional[int] = None,
|
||||
pad_token_id: Optional[int] = None,
|
||||
eos_token_id: Optional[int] = None,
|
||||
provided_constraints: List[ABCBloomConstraint] = [],
|
||||
) -> torch.LongTensor:
|
||||
"""
|
||||
Generates sequences of token ids for models with a language modeling head. Uses greedy search.
|
||||
|
||||
:param input_ids: The input tokens to the model.
|
||||
:param max_length: The maximum length of the sequence to generate.
|
||||
:param pad_token_id: The id of the padding token.
|
||||
:param eos_token_id: The id of the end of sentence token.
|
||||
:param provided_constraints: A list of constraints to use.
|
||||
"""
|
||||
return self.generate(
|
||||
inputs=input_ids,
|
||||
max_new_tokens=max_length,
|
||||
pad_token_id=pad_token_id,
|
||||
eos_token_id=eos_token_id,
|
||||
decoding_algorithm=GreedyAlgorithm(),
|
||||
provided_constraints=provided_constraints,
|
||||
)
|
||||
|
||||
def sample(
|
||||
self,
|
||||
input_ids: torch.LongTensor,
|
||||
temperature: float = 1.0,
|
||||
top_k: Optional[int] = None,
|
||||
top_p: Optional[float] = None,
|
||||
max_length: Optional[int] = None,
|
||||
pad_token_id: Optional[int] = None,
|
||||
eos_token_id: Optional[int] = None,
|
||||
provided_constraints: List[ABCBloomConstraint] = [],
|
||||
) -> torch.LongTensor:
|
||||
"""
|
||||
Generates sequences of token ids for models with a language modeling head. Uses multinomial sampling.
|
||||
If top_k is provided, uses top_k sampling. If top_p is provided, uses nucleus sampling.
|
||||
|
||||
:param: input_ids: The input tokens to the model.
|
||||
:param: temperature: The temperature to use for sampling.
|
||||
:param: top_k: The number of samples to use for top_k sampling.
|
||||
:param: top_p: The probability of using top_p sampling.
|
||||
:param: max_length: The maximum length of the sequence to generate.
|
||||
:param: pad_token_id: The id of the padding token.
|
||||
:param: eos_token_id: The id of the end of sentence token.
|
||||
:param: provided_constraints: A list of constraints to use.
|
||||
"""
|
||||
|
||||
return self.generate(
|
||||
inputs=input_ids,
|
||||
max_new_tokens=max_length,
|
||||
pad_token_id=pad_token_id,
|
||||
eos_token_id=eos_token_id,
|
||||
decoding_algorithm=self._choose_sample_algorithm(temperature, top_k, top_p),
|
||||
provided_constraints=provided_constraints,
|
||||
)
|
||||
|
||||
def beam_search(
|
||||
self,
|
||||
input_ids: torch.LongTensor,
|
||||
num_beams: int = 1,
|
||||
max_length: Optional[int] = None,
|
||||
pad_token_id: Optional[int] = None,
|
||||
eos_token_id: Optional[int] = None,
|
||||
provided_constraints: List[ABCBloomConstraint] = [],
|
||||
) -> torch.LongTensor:
|
||||
"""
|
||||
Generates sequences of token ids for models with a language modeling head. Uses beam search.
|
||||
|
||||
:param input_ids: The input tokens to the model.
|
||||
:param num_beams: The number of beams to use.
|
||||
:param max_length: The maximum length of the sequence to generate.
|
||||
:param pad_token_id: The id of the padding token.
|
||||
:param eos_token_id: The id of the end of sentence token.
|
||||
:param provided_constraints: A list of constraints to use.
|
||||
"""
|
||||
decoding_algorithm = BeamSearchAlgorithm(
|
||||
num_beams=num_beams,
|
||||
batch_size=input_ids.size(0),
|
||||
)
|
||||
return self.generate(
|
||||
inputs=input_ids,
|
||||
num_beams=num_beams,
|
||||
max_new_tokens=max_length,
|
||||
pad_token_id=pad_token_id,
|
||||
eos_token_id=eos_token_id,
|
||||
decoding_algorithm=decoding_algorithm,
|
||||
provided_constraints=provided_constraints,
|
||||
)
|
||||
|
||||
def beam_sample(
|
||||
self,
|
||||
input_ids: torch.LongTensor,
|
||||
max_length: Optional[int] = None,
|
||||
pad_token_id: Optional[int] = None,
|
||||
eos_token_id: Optional[int] = None,
|
||||
provided_constraints: List[ABCBloomConstraint] = [],
|
||||
) -> torch.LongTensor:
|
||||
raise NotImplementedError
|
||||
|
||||
def group_beam_search(
|
||||
self,
|
||||
input_ids: torch.LongTensor,
|
||||
max_length: Optional[int] = None,
|
||||
pad_token_id: Optional[int] = None,
|
||||
eos_token_id: Optional[int] = None,
|
||||
provided_constraints: List[ABCBloomConstraint] = [],
|
||||
) -> torch.LongTensor:
|
||||
raise NotImplementedError
|
||||
|
||||
def _choose_sample_algorithm(
|
||||
self,
|
||||
temperature: float = 1.0,
|
||||
top_k: Optional[int] = None,
|
||||
top_p: Optional[float] = None,
|
||||
) -> DecodingAlgorithm:
|
||||
if (top_k is not None) and (top_p is not None):
|
||||
raise ValueError("You have to provide only top_k or top_p for sampling")
|
||||
if top_k is not None:
|
||||
return TopKAlgorithm(top_k, temperature)
|
||||
elif top_p is not None:
|
||||
return NucleusAlgorithm(top_p, temperature)
|
||||
else:
|
||||
return SamplingAlgorithm(temperature)
|
||||
return kwargs
|
||||
|
||||
def _get_constraints(
|
||||
self,
|
||||
inputs: Optional[torch.Tensor] = None,
|
||||
eos_token_id: Optional[int] = None,
|
||||
pad_token_id: Optional[int] = None,
|
||||
provided_constraints: List[ABCBloomConstraint] = [],
|
||||
) -> List[ABCBloomConstraint]:
|
||||
constraints = []
|
||||
constraints.extend(provided_constraints)
|
||||
constraints.append(EosConstraint(inputs, eos_token_id, pad_token_id))
|
||||
return constraints
|
||||
@staticmethod
|
||||
def _reorder_cache(past_key_values: RemotePastKeyValues, beam_idx: torch.LongTensor) -> RemotePastKeyValues:
|
||||
return dataclasses.replace(past_key_values, hypo_ids=beam_idx)
|
||||
|
||||
@ -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,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
|
||||
@ -1,128 +0,0 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
|
||||
TokenIds = torch.Tensor
|
||||
HypoIds = torch.Tensor
|
||||
|
||||
|
||||
class DecodingAlgorithm(ABC):
|
||||
"""
|
||||
An abstract class for decoding algorithms. Describes the base function of those algorithms:
|
||||
they have to select new tokens and provide the corresponding hypotheses.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def __call__(self, logits: torch.Tensor) -> Tuple[TokenIds, HypoIds]:
|
||||
"""
|
||||
:param logits: A tensor of shape (batch_size, seq_length, vocab_size)
|
||||
:return: A tuple of selected token ids and corresponding hypotheses.
|
||||
The shape of the token ids is (batch_size, seq_length), and the shape of the hypotheses is (batch_size)
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
class GreedyAlgorithm(DecodingAlgorithm):
|
||||
"""
|
||||
The simplest algorithm for decoding. It selects the most probable token.
|
||||
"""
|
||||
|
||||
def __call__(self, logits: torch.Tensor) -> Tuple[TokenIds, HypoIds]:
|
||||
"""
|
||||
Returns the most probable token. The second returned object is always a range of integers
|
||||
from 0 to batch_size - 1.
|
||||
"""
|
||||
return logits.max(-1)[1].unsqueeze(1), torch.arange(logits.size(0))
|
||||
|
||||
|
||||
class SamplingAlgorithm(DecodingAlgorithm):
|
||||
def __init__(self, temperature: float = 1.0):
|
||||
self.temperature = temperature
|
||||
|
||||
def sample(self, logits: torch.Tensor, indices_to_remove: torch.Tensor) -> Tuple[TokenIds, HypoIds]:
|
||||
"""
|
||||
:param logits: A tensor of shape (batch_size * num_hypos, vocab_size)
|
||||
:param indices_to_remove: A bool tensor of shape (batch_size * num_hypos, vocab_size)
|
||||
:return: A tuple of selected token ids and corresponding hypotheses.
|
||||
The shape of the token ids is (batch_size, seq_length), and the shape of the hypotheses is (batch_size).
|
||||
"""
|
||||
logits[indices_to_remove] = -float("Inf")
|
||||
probs = torch.softmax(logits / self.temperature, -1)
|
||||
return torch.multinomial(probs, num_samples=1), torch.arange(logits.size(0))
|
||||
|
||||
def __call__(self, logits: torch.Tensor) -> Tuple[TokenIds, HypoIds]:
|
||||
indices_to_remove = torch.full_like(logits, False, dtype=torch.bool)
|
||||
return self.sample(logits, indices_to_remove)
|
||||
|
||||
|
||||
class TopKAlgorithm(SamplingAlgorithm):
|
||||
def __init__(self, top_k: int, temperature: float = 1.0) -> None:
|
||||
super().__init__(temperature=temperature)
|
||||
self.top_k = top_k
|
||||
|
||||
def __call__(self, logits: torch.Tensor) -> Tuple[TokenIds, HypoIds]:
|
||||
indices_to_remove = logits < torch.topk(logits, self.top_k, dim=-1)[0][..., -1, None]
|
||||
return self.sample(logits, indices_to_remove)
|
||||
|
||||
|
||||
class NucleusAlgorithm(SamplingAlgorithm):
|
||||
def __init__(self, top_p: float, temperature: float = 1.0) -> None:
|
||||
super().__init__(temperature=temperature)
|
||||
self.top_p = top_p
|
||||
|
||||
def __call__(self, logits: torch.Tensor) -> Tuple[TokenIds, HypoIds]:
|
||||
sorted_logits, sorted_indices = torch.sort(logits, descending=False, dim=-1)
|
||||
probs = torch.softmax(sorted_logits / self.temperature, -1)
|
||||
cumulative_probs = torch.cumsum(probs, dim=-1)
|
||||
|
||||
sorted_indices_to_remove = cumulative_probs <= (1 - self.top_p)
|
||||
|
||||
indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
|
||||
return self.sample(logits, indices_to_remove)
|
||||
|
||||
|
||||
class BeamSearchAlgorithm(DecodingAlgorithm):
|
||||
def __init__(self, num_beams: int, batch_size: int) -> None:
|
||||
self.num_beams = num_beams
|
||||
self.batch_size = batch_size
|
||||
|
||||
self._batch_beams = [list() for _ in range(batch_size)]
|
||||
|
||||
def __call__(self, logits: torch.Tensor):
|
||||
sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
|
||||
probs = torch.log_softmax(sorted_logits, -1)
|
||||
|
||||
if len(self._batch_beams[0]) > 0:
|
||||
for batch_idx in range(self.batch_size):
|
||||
new_beams = []
|
||||
cur_beams = self._batch_beams[batch_idx]
|
||||
for beam_idx in range(len(cur_beams)):
|
||||
probs_idx = batch_idx + beam_idx * self.batch_size
|
||||
new_beam = cur_beams[beam_idx]
|
||||
for hypo_idx in range(self.num_beams):
|
||||
new_beams.append(
|
||||
(new_beam[0] + probs[probs_idx, hypo_idx].item(), beam_idx * self.num_beams + hypo_idx)
|
||||
)
|
||||
self._batch_beams[batch_idx] = sorted(new_beams, reverse=True)[: self.num_beams]
|
||||
else:
|
||||
for batch_idx in range(self.batch_size):
|
||||
for beam_idx in range(self.num_beams):
|
||||
self._batch_beams[batch_idx].append((probs[batch_idx, beam_idx].item(), beam_idx))
|
||||
|
||||
return_hypos = []
|
||||
return_tokens = []
|
||||
for batch_idx in range(self.batch_size):
|
||||
cur_beam = self._batch_beams[batch_idx]
|
||||
return_hypos.append(list())
|
||||
return_tokens.append(list())
|
||||
for beam in cur_beam:
|
||||
beam_idx = beam[1] // self.num_beams
|
||||
hypo_idx = batch_idx + beam_idx * self.batch_size
|
||||
token_idx = beam[1] % self.num_beams
|
||||
return_hypos[-1].append(hypo_idx)
|
||||
return_tokens[-1].append([sorted_indices[hypo_idx, token_idx].item()])
|
||||
return_hypos = [hypo_idx for hypo_indexes in zip(*return_hypos) for hypo_idx in hypo_indexes]
|
||||
return_tokens = [token_idx for token_indexes in zip(*return_tokens) for token_idx in token_indexes]
|
||||
|
||||
return torch.tensor(return_tokens), torch.tensor(return_hypos)
|
||||
@ -1,51 +0,0 @@
|
||||
from abc import ABC
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
class ABCBloomConstraint(ABC):
|
||||
"""
|
||||
Base class of all kind of decoding constraints. It can be used to implement a new constraint.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
|
||||
def __call__(self, tokens_id: torch.Tensor, logits: torch.Tensor, hypo_ids: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
This method is called by the decoding algorithm to apply the constraint. It changes and returns new logits.
|
||||
:param tokens_id: The token id of the last chosen token.
|
||||
:param logits: The logits from the Bloom model.
|
||||
:param hypo_ids: The hypothesis ids of the last tokens.
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
class EosConstraint(ABCBloomConstraint):
|
||||
"""
|
||||
This constrained repeats EOS token if it was generated on the previous step.
|
||||
Args:
|
||||
prefix: The prefix of the sequence.
|
||||
eos_token_id: The id of the end of sentence token.
|
||||
pad_token_id: The id of the padding token.
|
||||
min_logits: The minimum logits that can be generated. Default: -1e6.
|
||||
"""
|
||||
|
||||
def __init__(self, prefix: torch.Tensor, eos_token_id: int, pad_token_id: int, min_logits: float = -1e8) -> None:
|
||||
self.eos_token_id = eos_token_id
|
||||
self.min_logits = min_logits
|
||||
self.past_tokens = None
|
||||
|
||||
self.wait_until_starting = (prefix == pad_token_id).sum(1).unsqueeze(1)
|
||||
|
||||
def __call__(self, tokens_id: torch.Tensor, logits: torch.Tensor, hypo_ids: torch.Tensor) -> torch.Tensor:
|
||||
if self.past_tokens is not None:
|
||||
mask = (self.wait_until_starting < 0) & (self.past_tokens == self.eos_token_id)
|
||||
logits += self.min_logits * mask
|
||||
logits[mask[:, 0], self.eos_token_id] = 0
|
||||
|
||||
if tokens_id is not None:
|
||||
self.past_tokens = tokens_id
|
||||
self.wait_until_starting -= 1
|
||||
|
||||
return logits
|
||||
@ -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,
|
||||
)
|
||||
Loading…
Reference in New Issue