evelynmitchell
1 year ago
committed by
GitHub
commit
a3fce76d0e
@ -0,0 +1,82 @@
|
||||
import pkgutil
|
||||
import inspect
|
||||
|
||||
|
||||
class ModelRegistry:
|
||||
"""
|
||||
A registry for storing and querying models.
|
||||
|
||||
Attributes:
|
||||
models (dict): A dictionary of model names and corresponding model classes.
|
||||
|
||||
Methods:
|
||||
__init__(): Initializes the ModelRegistry object and retrieves all available models.
|
||||
_get_all_models(): Retrieves all available models from the models package.
|
||||
query(text): Queries the models based on the given text and returns a dictionary of matching models.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.models = self._get_all_models()
|
||||
|
||||
def _get_all_models(self):
|
||||
"""
|
||||
Retrieves all available models from the models package.
|
||||
|
||||
Returns:
|
||||
dict: A dictionary of model names and corresponding model classes.
|
||||
"""
|
||||
models = {}
|
||||
for importer, modname, ispkg in pkgutil.iter_modules(
|
||||
models.__path__
|
||||
):
|
||||
module = importer.find_module(modname).load_module(
|
||||
modname
|
||||
)
|
||||
for name, obj in inspect.getmembers(module):
|
||||
if inspect.isclass(obj):
|
||||
models[name] = obj
|
||||
return models
|
||||
|
||||
def query(self, text):
|
||||
"""
|
||||
Queries the models based on the given text and returns a dictionary of matching models.
|
||||
|
||||
Args:
|
||||
text (str): The text to search for in the model names.
|
||||
|
||||
Returns:
|
||||
dict: A dictionary of matching model names and corresponding model classes.
|
||||
"""
|
||||
return {
|
||||
name: model
|
||||
for name, model in self.models.items()
|
||||
if text in name
|
||||
}
|
||||
|
||||
def run_model(
|
||||
self, model_name: str, task: str, img: str, *args, **kwargs
|
||||
):
|
||||
"""
|
||||
Runs the specified model for the given task and image.
|
||||
|
||||
Args:
|
||||
model_name (str): The name of the model to run.
|
||||
task (str): The task to perform using the model.
|
||||
img (str): The image to process.
|
||||
*args: Additional positional arguments to pass to the model's run method.
|
||||
**kwargs: Additional keyword arguments to pass to the model's run method.
|
||||
|
||||
Returns:
|
||||
The result of running the model.
|
||||
|
||||
Raises:
|
||||
ValueError: If the specified model is not found in the model registry.
|
||||
"""
|
||||
if model_name not in self.models:
|
||||
raise ValueError(f"Model {model_name} not found")
|
||||
|
||||
# Get the model
|
||||
model = self.models[model_name]
|
||||
|
||||
# Run the model
|
||||
return model.run(task, img, *args, **kwargs)
|
||||
@ -0,0 +1,299 @@
|
||||
"""Sampling parameters for text generation."""
|
||||
from enum import IntEnum
|
||||
from functools import cached_property
|
||||
from typing import Callable, List, Optional, Union
|
||||
|
||||
import torch
|
||||
|
||||
_SAMPLING_EPS = 1e-5
|
||||
|
||||
|
||||
class SamplingType(IntEnum):
|
||||
GREEDY = 0
|
||||
RANDOM = 1
|
||||
BEAM = 2
|
||||
|
||||
|
||||
LogitsProcessor = Callable[[List[int], torch.Tensor], torch.Tensor]
|
||||
"""LogitsProcessor is a function that takes a list of previously generated
|
||||
tokens and a tensor of the logits for the next token, and returns a modified
|
||||
tensor of logits to sample from."""
|
||||
|
||||
|
||||
class SamplingParams:
|
||||
"""Sampling parameters for text generation.
|
||||
|
||||
Overall, we follow the sampling parameters from the OpenAI text completion
|
||||
API (https://platform.openai.com/docs/api-reference/completions/create).
|
||||
In addition, we support beam search, which is not supported by OpenAI.
|
||||
|
||||
Args:
|
||||
n: Number of output sequences to return for the given prompt.
|
||||
best_of: Number of output sequences that are generated from the prompt.
|
||||
From these `best_of` sequences, the top `n` sequences are returned.
|
||||
`best_of` must be greater than or equal to `n`. This is treated as
|
||||
the beam width when `use_beam_search` is True. By default, `best_of`
|
||||
is set to `n`.
|
||||
presence_penalty: Float that penalizes new tokens based on whether they
|
||||
appear in the generated text so far. Values > 0 encourage the model
|
||||
to use new tokens, while values < 0 encourage the model to repeat
|
||||
tokens.
|
||||
frequency_penalty: Float that penalizes new tokens based on their
|
||||
frequency in the generated text so far. Values > 0 encourage the
|
||||
model to use new tokens, while values < 0 encourage the model to
|
||||
repeat tokens.
|
||||
repetition_penalty: Float that penalizes new tokens based on whether
|
||||
they appear in the prompt and the generated text so far. Values > 1
|
||||
encourage the model to use new tokens, while values < 1 encourage
|
||||
the model to repeat tokens.
|
||||
temperature: Float that controls the randomness of the sampling. Lower
|
||||
values make the model more deterministic, while higher values make
|
||||
the model more random. Zero means greedy sampling.
|
||||
top_p: Float that controls the cumulative probability of the top tokens
|
||||
to consider. Must be in (0, 1]. Set to 1 to consider all tokens.
|
||||
top_k: Integer that controls the number of top tokens to consider. Set
|
||||
to -1 to consider all tokens.
|
||||
min_p: Float that represents the minimum probability for a token to be
|
||||
considered, relative to the probability of the most likely token.
|
||||
Must be in [0, 1]. Set to 0 to disable this.
|
||||
use_beam_search: Whether to use beam search instead of sampling.
|
||||
length_penalty: Float that penalizes sequences based on their length.
|
||||
Used in beam search.
|
||||
early_stopping: Controls the stopping condition for beam search. It
|
||||
accepts the following values: `True`, where the generation stops as
|
||||
soon as there are `best_of` complete candidates; `False`, where an
|
||||
heuristic is applied and the generation stops when is it very
|
||||
unlikely to find better candidates; `"never"`, where the beam search
|
||||
procedure only stops when there cannot be better candidates
|
||||
(canonical beam search algorithm).
|
||||
stop: List of strings that stop the generation when they are generated.
|
||||
The returned output will not contain the stop strings.
|
||||
stop_token_ids: List of tokens that stop the generation when they are
|
||||
generated. The returned output will contain the stop tokens unless
|
||||
the stop tokens are special tokens.
|
||||
include_stop_str_in_output: Whether to include the stop strings in output
|
||||
text. Defaults to False.
|
||||
ignore_eos: Whether to ignore the EOS token and continue generating
|
||||
tokens after the EOS token is generated.
|
||||
max_tokens: Maximum number of tokens to generate per output sequence.
|
||||
logprobs: Number of log probabilities to return per output token.
|
||||
Note that the implementation follows the OpenAI API: The return
|
||||
result includes the log probabilities on the `logprobs` most likely
|
||||
tokens, as well the chosen tokens. The API will always return the
|
||||
log probability of the sampled token, so there may be up to
|
||||
`logprobs+1` elements in the response.
|
||||
prompt_logprobs: Number of log probabilities to return per prompt token.
|
||||
skip_special_tokens: Whether to skip special tokens in the output.
|
||||
spaces_between_special_tokens: Whether to add spaces between special
|
||||
tokens in the output. Defaults to True.
|
||||
logits_processors: List of functions that modify logits based on
|
||||
previously generated tokens.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
n: int = 1,
|
||||
best_of: Optional[int] = None,
|
||||
presence_penalty: float = 0.0,
|
||||
frequency_penalty: float = 0.0,
|
||||
repetition_penalty: float = 1.0,
|
||||
temperature: float = 1.0,
|
||||
top_p: float = 1.0,
|
||||
top_k: int = -1,
|
||||
min_p: float = 0.0,
|
||||
use_beam_search: bool = False,
|
||||
length_penalty: float = 1.0,
|
||||
early_stopping: Union[bool, str] = False,
|
||||
stop: Optional[Union[str, List[str]]] = None,
|
||||
stop_token_ids: Optional[List[int]] = None,
|
||||
include_stop_str_in_output: bool = False,
|
||||
ignore_eos: bool = False,
|
||||
max_tokens: Optional[int] = 16,
|
||||
logprobs: Optional[int] = None,
|
||||
prompt_logprobs: Optional[int] = None,
|
||||
skip_special_tokens: bool = True,
|
||||
spaces_between_special_tokens: bool = True,
|
||||
logits_processors: Optional[List[LogitsProcessor]] = None,
|
||||
) -> None:
|
||||
self.n = n
|
||||
self.best_of = best_of if best_of is not None else n
|
||||
self.presence_penalty = presence_penalty
|
||||
self.frequency_penalty = frequency_penalty
|
||||
self.repetition_penalty = repetition_penalty
|
||||
self.temperature = temperature
|
||||
self.top_p = top_p
|
||||
self.top_k = top_k
|
||||
self.min_p = min_p
|
||||
self.use_beam_search = use_beam_search
|
||||
self.length_penalty = length_penalty
|
||||
self.early_stopping = early_stopping
|
||||
if stop is None:
|
||||
self.stop = []
|
||||
elif isinstance(stop, str):
|
||||
self.stop = [stop]
|
||||
else:
|
||||
self.stop = list(stop)
|
||||
if stop_token_ids is None:
|
||||
self.stop_token_ids = []
|
||||
else:
|
||||
self.stop_token_ids = list(stop_token_ids)
|
||||
self.ignore_eos = ignore_eos
|
||||
self.max_tokens = max_tokens
|
||||
self.logprobs = logprobs
|
||||
self.prompt_logprobs = prompt_logprobs
|
||||
self.skip_special_tokens = skip_special_tokens
|
||||
self.spaces_between_special_tokens = (
|
||||
spaces_between_special_tokens
|
||||
)
|
||||
self.logits_processors = logits_processors
|
||||
self.include_stop_str_in_output = include_stop_str_in_output
|
||||
self._verify_args()
|
||||
if self.use_beam_search:
|
||||
self._verify_beam_search()
|
||||
else:
|
||||
self._verify_non_beam_search()
|
||||
if self.temperature < _SAMPLING_EPS:
|
||||
# Zero temperature means greedy sampling.
|
||||
self.top_p = 1.0
|
||||
self.top_k = -1
|
||||
self.min_p = 0.0
|
||||
self._verify_greedy_sampling()
|
||||
|
||||
def _verify_args(self) -> None:
|
||||
if self.n < 1:
|
||||
raise ValueError(f"n must be at least 1, got {self.n}.")
|
||||
if self.best_of < self.n:
|
||||
raise ValueError(
|
||||
"best_of must be greater than or equal to n, "
|
||||
f"got n={self.n} and best_of={self.best_of}."
|
||||
)
|
||||
if not -2.0 <= self.presence_penalty <= 2.0:
|
||||
raise ValueError(
|
||||
"presence_penalty must be in [-2, 2], got "
|
||||
f"{self.presence_penalty}."
|
||||
)
|
||||
if not -2.0 <= self.frequency_penalty <= 2.0:
|
||||
raise ValueError(
|
||||
"frequency_penalty must be in [-2, 2], got "
|
||||
f"{self.frequency_penalty}."
|
||||
)
|
||||
if not 0.0 < self.repetition_penalty <= 2.0:
|
||||
raise ValueError(
|
||||
"repetition_penalty must be in (0, 2], got "
|
||||
f"{self.repetition_penalty}."
|
||||
)
|
||||
if self.temperature < 0.0:
|
||||
raise ValueError(
|
||||
"temperature must be non-negative, got"
|
||||
f" {self.temperature}."
|
||||
)
|
||||
if not 0.0 < self.top_p <= 1.0:
|
||||
raise ValueError(
|
||||
f"top_p must be in (0, 1], got {self.top_p}."
|
||||
)
|
||||
if self.top_k < -1 or self.top_k == 0:
|
||||
raise ValueError(
|
||||
"top_k must be -1 (disable), or at least 1, "
|
||||
f"got {self.top_k}."
|
||||
)
|
||||
if not 0.0 <= self.min_p <= 1.0:
|
||||
raise ValueError(
|
||||
f"min_p must be in [0, 1], got {self.min_p}."
|
||||
)
|
||||
if self.max_tokens is not None and self.max_tokens < 1:
|
||||
raise ValueError(
|
||||
"max_tokens must be at least 1, got"
|
||||
f" {self.max_tokens}."
|
||||
)
|
||||
if self.logprobs is not None and self.logprobs < 0:
|
||||
raise ValueError(
|
||||
f"logprobs must be non-negative, got {self.logprobs}."
|
||||
)
|
||||
if (
|
||||
self.prompt_logprobs is not None
|
||||
and self.prompt_logprobs < 0
|
||||
):
|
||||
raise ValueError(
|
||||
"prompt_logprobs must be non-negative, got "
|
||||
f"{self.prompt_logprobs}."
|
||||
)
|
||||
|
||||
def _verify_beam_search(self) -> None:
|
||||
if self.best_of == 1:
|
||||
raise ValueError(
|
||||
"best_of must be greater than 1 when using beam "
|
||||
f"search. Got {self.best_of}."
|
||||
)
|
||||
if self.temperature > _SAMPLING_EPS:
|
||||
raise ValueError(
|
||||
"temperature must be 0 when using beam search."
|
||||
)
|
||||
if self.top_p < 1.0 - _SAMPLING_EPS:
|
||||
raise ValueError(
|
||||
"top_p must be 1 when using beam search."
|
||||
)
|
||||
if self.top_k != -1:
|
||||
raise ValueError(
|
||||
"top_k must be -1 when using beam search."
|
||||
)
|
||||
if self.early_stopping not in [True, False, "never"]:
|
||||
raise ValueError(
|
||||
"early_stopping must be True, False, or 'never', "
|
||||
f"got {self.early_stopping}."
|
||||
)
|
||||
|
||||
def _verify_non_beam_search(self) -> None:
|
||||
if self.early_stopping is not False:
|
||||
raise ValueError(
|
||||
"early_stopping is not effective and must be "
|
||||
"False when not using beam search."
|
||||
)
|
||||
if (
|
||||
self.length_penalty < 1.0 - _SAMPLING_EPS
|
||||
or self.length_penalty > 1.0 + _SAMPLING_EPS
|
||||
):
|
||||
raise ValueError(
|
||||
"length_penalty is not effective and must be the "
|
||||
"default value of 1.0 when not using beam search."
|
||||
)
|
||||
|
||||
def _verify_greedy_sampling(self) -> None:
|
||||
if self.best_of > 1:
|
||||
raise ValueError(
|
||||
"best_of must be 1 when using greedy sampling."
|
||||
f"Got {self.best_of}."
|
||||
)
|
||||
|
||||
@cached_property
|
||||
def sampling_type(self) -> SamplingType:
|
||||
if self.use_beam_search:
|
||||
return SamplingType.BEAM
|
||||
if self.temperature < _SAMPLING_EPS:
|
||||
return SamplingType.GREEDY
|
||||
return SamplingType.RANDOM
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"SamplingParams(n={self.n}, "
|
||||
f"best_of={self.best_of}, "
|
||||
f"presence_penalty={self.presence_penalty}, "
|
||||
f"frequency_penalty={self.frequency_penalty}, "
|
||||
f"repetition_penalty={self.repetition_penalty}, "
|
||||
f"temperature={self.temperature}, "
|
||||
f"top_p={self.top_p}, "
|
||||
f"top_k={self.top_k}, "
|
||||
f"min_p={self.min_p}, "
|
||||
f"use_beam_search={self.use_beam_search}, "
|
||||
f"length_penalty={self.length_penalty}, "
|
||||
f"early_stopping={self.early_stopping}, "
|
||||
f"stop={self.stop}, "
|
||||
f"stop_token_ids={self.stop_token_ids}, "
|
||||
f"include_stop_str_in_output={self.include_stop_str_in_output}, "
|
||||
f"ignore_eos={self.ignore_eos}, "
|
||||
f"max_tokens={self.max_tokens}, "
|
||||
f"logprobs={self.logprobs}, "
|
||||
f"prompt_logprobs={self.prompt_logprobs}, "
|
||||
f"skip_special_tokens={self.skip_special_tokens}, "
|
||||
"spaces_between_special_tokens="
|
||||
f"{self.spaces_between_special_tokens})"
|
||||
)
|
||||
@ -0,0 +1,80 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import torch
|
||||
|
||||
|
||||
def continuous_tensor(
|
||||
inputs: torch.Tensor, seq_length: torch.LongTensor
|
||||
):
|
||||
"""Convert batched tensor to continuous tensor.
|
||||
|
||||
Args:
|
||||
inputs (Tensor): batched tensor.
|
||||
seq_length (Tensor): length of each sequence.
|
||||
|
||||
Return:
|
||||
Tensor: continuoused tensor.
|
||||
"""
|
||||
assert inputs.dim() > 1
|
||||
if inputs.size(1) == 1:
|
||||
return inputs.reshape(1, -1)
|
||||
|
||||
inputs = [inp[:slen] for inp, slen in zip(inputs, seq_length)]
|
||||
|
||||
inputs = torch.cat(inputs).unsqueeze(0)
|
||||
return inputs
|
||||
|
||||
|
||||
def batch_tensor(inputs: torch.Tensor, seq_length: torch.LongTensor):
|
||||
"""Convert continuoused tensor to batched tensor.
|
||||
|
||||
Args:
|
||||
inputs (Tensor): continuoused tensor.
|
||||
seq_length (Tensor): length of each sequence.
|
||||
|
||||
Return:
|
||||
Tensor: batched tensor.
|
||||
"""
|
||||
from torch.nn.utils.rnn import pad_sequence
|
||||
|
||||
end_loc = seq_length.cumsum(0)
|
||||
start_loc = end_loc - seq_length
|
||||
|
||||
inputs = [
|
||||
inputs[0, sloc:eloc] for sloc, eloc in zip(start_loc, end_loc)
|
||||
]
|
||||
inputs = pad_sequence(inputs, batch_first=True)
|
||||
return inputs
|
||||
|
||||
|
||||
def page_cache(
|
||||
paged_cache: torch.Tensor,
|
||||
batched_cache: torch.Tensor,
|
||||
cache_length: torch.Tensor,
|
||||
block_offsets: torch.Tensor,
|
||||
permute_head: bool = True,
|
||||
):
|
||||
"""Convert batched cache to paged cache.
|
||||
|
||||
Args:
|
||||
paged_cache (Tensor): Output paged cache.
|
||||
batched_cache (Tensor): Input batched cache.
|
||||
cache_length (Tensor): length of the cache.
|
||||
block_offsets (Tensor): Offset of each blocks.
|
||||
"""
|
||||
assert block_offsets.dim() == 2
|
||||
block_size = paged_cache.size(1)
|
||||
batch_size = batched_cache.size(0)
|
||||
if permute_head:
|
||||
batched_cache = batched_cache.permute(0, 2, 1, 3)
|
||||
|
||||
for b_idx in range(batch_size):
|
||||
cache_len = cache_length[b_idx]
|
||||
b_cache = batched_cache[b_idx]
|
||||
block_off = block_offsets[b_idx]
|
||||
block_off_idx = 0
|
||||
for s_start in range(0, cache_len, block_size):
|
||||
s_end = min(s_start + block_size, cache_len)
|
||||
s_len = s_end - s_start
|
||||
b_off = block_off[block_off_idx]
|
||||
paged_cache[b_off, :s_len] = b_cache[s_start:s_end]
|
||||
block_off_idx += 1
|
||||
@ -0,0 +1,315 @@
|
||||
from typing import Callable, Union
|
||||
|
||||
import torch
|
||||
from torch import Tensor, nn
|
||||
from torch.distributed._tensor import (
|
||||
DeviceMesh,
|
||||
DTensor,
|
||||
Replicate,
|
||||
Shard,
|
||||
distribute_tensor,
|
||||
)
|
||||
from zeta.nn import QuantizedLN
|
||||
|
||||
|
||||
try:
|
||||
from peft.tuners.lora import Linear as LoRALinear
|
||||
except ImportError:
|
||||
|
||||
class LoRALinear:
|
||||
pass
|
||||
|
||||
|
||||
def try_to_local(tensor: Union[Tensor, DTensor]):
|
||||
"""Try to convert DTensor to Tensor.
|
||||
|
||||
Args:
|
||||
tensor (Tensor|DTensor): Tensor to convert.
|
||||
"""
|
||||
if isinstance(tensor, DTensor):
|
||||
tensor = tensor.to_local()
|
||||
return tensor
|
||||
|
||||
|
||||
def module_to_local(module: nn.Module):
|
||||
"""convert all DTensor parameters to Tensor parameters in module.
|
||||
|
||||
Args:
|
||||
module (Module): Module to convert.
|
||||
"""
|
||||
for name, mod in module.named_children():
|
||||
module_to_local(mod)
|
||||
|
||||
for name, param in module.named_parameters(recurse=False):
|
||||
module.register_parameter(
|
||||
name, nn.Parameter(try_to_local(param))
|
||||
)
|
||||
|
||||
for name, buf in module.named_buffers(recurse=False):
|
||||
module.register_buffer(name, try_to_local(buf))
|
||||
|
||||
|
||||
def rowwise_parallelize_linear(
|
||||
module: nn.Module, device_mesh: DeviceMesh, to_local: bool = False
|
||||
) -> None:
|
||||
"""
|
||||
This function parallelizes the input :class:`nn.Linear` module in
|
||||
:class:`RowwiseParallel` style.
|
||||
|
||||
Args:
|
||||
module (:class:`nn.Module`):
|
||||
The :class:`nn.Linear` module to be parallelized.
|
||||
device_mesh (:class:`DeviceMesh`):
|
||||
Object which describes the mesh topology of devices.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
for name, param in module.named_parameters():
|
||||
dist_spec = (
|
||||
[Shard(1)]
|
||||
if name == "weight"
|
||||
else [Replicate()] # type: ignore[list-item]
|
||||
)
|
||||
|
||||
dist_tensor = distribute_tensor(param, device_mesh, dist_spec)
|
||||
if to_local:
|
||||
dist_tensor = try_to_local(dist_tensor)
|
||||
if name == "bias":
|
||||
# rowwise linear would add bias more than ones.
|
||||
dist_tensor /= device_mesh.size()
|
||||
dist_param = torch.nn.Parameter(dist_tensor)
|
||||
module.register_parameter(name, dist_param)
|
||||
|
||||
# Weight, bias and scale are registered as buffer in QLinear
|
||||
for name, buffer in module.named_buffers():
|
||||
dist_spec = (
|
||||
[Shard(1)]
|
||||
if name == "weight"
|
||||
else [Replicate()] # type: ignore[list-item]
|
||||
)
|
||||
|
||||
dist_tensor = distribute_tensor(
|
||||
buffer, device_mesh, dist_spec
|
||||
)
|
||||
if to_local:
|
||||
dist_tensor = try_to_local(dist_tensor)
|
||||
if name == "bias":
|
||||
# rowwise linear would add bias more than ones.
|
||||
dist_tensor /= device_mesh.size()
|
||||
module.register_buffer(name, dist_tensor)
|
||||
|
||||
dist_tensor = distribute_tensor(
|
||||
buffer, device_mesh, dist_spec
|
||||
)
|
||||
if to_local:
|
||||
dist_tensor = try_to_local(dist_tensor)
|
||||
module.register_buffer(name, dist_tensor)
|
||||
|
||||
|
||||
def rowwise_parallelize_loralinear(
|
||||
module: LoRALinear,
|
||||
device_mesh: DeviceMesh,
|
||||
to_local: bool = False,
|
||||
) -> None:
|
||||
"""rowwize parallelize lora linear.
|
||||
|
||||
Read S-LoRA for more detail.
|
||||
"""
|
||||
rowwise_parallelize_linear(
|
||||
module.base_layer, device_mesh=device_mesh, to_local=to_local
|
||||
)
|
||||
for mod in module.lora_A.values():
|
||||
rowwise_parallelize_linear(
|
||||
mod, device_mesh=device_mesh, to_local=to_local
|
||||
)
|
||||
for mod in module.lora_B.values():
|
||||
colwise_parallelize_linear(
|
||||
mod, device_mesh=device_mesh, to_local=to_local
|
||||
)
|
||||
module._tp_mode = "rowwise"
|
||||
|
||||
|
||||
def rowwise_parallelize_linear_fn(
|
||||
module: nn.Module, device_mesh: DeviceMesh, to_local: bool = False
|
||||
) -> None:
|
||||
"""
|
||||
This function parallelizes the input :Linear module in
|
||||
:class:`RowwiseParallel` style.
|
||||
|
||||
Args:
|
||||
module (:class:`nn.Module`):
|
||||
The :class:`nn.Linear` module to be parallelized.
|
||||
device_mesh (:class:`DeviceMesh`):
|
||||
Object which describes the mesh topology of devices.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
if isinstance(module, (torch.nn.Linear, QuantizedLN)):
|
||||
return rowwise_parallelize_linear(
|
||||
module, device_mesh=device_mesh, to_local=to_local
|
||||
)
|
||||
elif isinstance(module, LoRALinear):
|
||||
return rowwise_parallelize_loralinear(
|
||||
module, device_mesh=device_mesh, to_local=to_local
|
||||
)
|
||||
else:
|
||||
raise TypeError(f"Unsupported module: {type(module)}")
|
||||
|
||||
|
||||
def colwise_parallelize_linear(
|
||||
module: nn.Module, device_mesh: DeviceMesh, to_local: bool = False
|
||||
) -> None:
|
||||
"""
|
||||
This function parallelizes the input :class:`nn.Linear` module in
|
||||
:class:`ColwiseParallel` style.
|
||||
|
||||
Args:
|
||||
module (:class:`nn.Module`):
|
||||
The :class:`nn.Linear` module to be parallelized.
|
||||
device_mesh (:class:`DeviceMesh`):
|
||||
Object which describes the mesh topology of devices.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
|
||||
for name, param in module.named_parameters():
|
||||
dist_tensor = distribute_tensor(
|
||||
param, device_mesh, [Shard(0)]
|
||||
)
|
||||
if to_local:
|
||||
dist_tensor = try_to_local(dist_tensor)
|
||||
dist_param = torch.nn.Parameter(dist_tensor)
|
||||
module.register_parameter(name, dist_param)
|
||||
# Weight, bias and scale are registered as buffer in QLinear
|
||||
for name, buffer in module.named_buffers():
|
||||
dist_tensor = distribute_tensor(
|
||||
buffer, device_mesh, [Shard(0)]
|
||||
)
|
||||
if to_local:
|
||||
dist_tensor = try_to_local(dist_tensor)
|
||||
module.register_buffer(name, dist_tensor)
|
||||
|
||||
|
||||
def colwise_parallelize_loralinear(
|
||||
module: nn.Module, device_mesh: DeviceMesh, to_local: bool = False
|
||||
) -> None:
|
||||
"""colwise parallelize lora linear."""
|
||||
colwise_parallelize_linear(
|
||||
module.base_layer, device_mesh=device_mesh, to_local=to_local
|
||||
)
|
||||
for mod in module.lora_A.values():
|
||||
colwise_parallelize_linear(
|
||||
mod, device_mesh=device_mesh, to_local=to_local
|
||||
)
|
||||
for mod in module.lora_B.values():
|
||||
colwise_parallelize_linear(
|
||||
mod, device_mesh=device_mesh, to_local=to_local
|
||||
)
|
||||
module._tp_mode = "colwise"
|
||||
|
||||
|
||||
def colwise_parallelize_linear_fn(
|
||||
module: nn.Module, device_mesh: DeviceMesh, to_local: bool = False
|
||||
) -> None:
|
||||
"""
|
||||
This function parallelizes the input :Linear module in
|
||||
:class:`ColwiseParallel` style.
|
||||
|
||||
Args:
|
||||
module (:class:`nn.Module`):
|
||||
The :class:`nn.Linear` module to be parallelized.
|
||||
device_mesh (:class:`DeviceMesh`):
|
||||
Object which describes the mesh topology of devices.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
if isinstance(module, (torch.nn.Linear, QuantizedLN)):
|
||||
return colwise_parallelize_linear(
|
||||
module, device_mesh=device_mesh, to_local=to_local
|
||||
)
|
||||
elif isinstance(module, LoRALinear):
|
||||
return colwise_parallelize_loralinear(
|
||||
module, device_mesh=device_mesh, to_local=to_local
|
||||
)
|
||||
else:
|
||||
raise TypeError(f"Unsupported module: {type(module)}")
|
||||
|
||||
|
||||
def _partition_module(
|
||||
mod_name: str,
|
||||
prefix: str,
|
||||
module: nn.Module,
|
||||
device_mesh: DeviceMesh,
|
||||
func: Callable,
|
||||
):
|
||||
"""partition module.
|
||||
|
||||
Parameters in module won't be force Replicated.
|
||||
|
||||
Args:
|
||||
mod_name (str): module name.
|
||||
prefix (str): Parameter prefix.
|
||||
module (Module): Module to be partitioned.
|
||||
device_mesh (DeviceMesh): The device mesh.
|
||||
func (Callable): partition callback
|
||||
"""
|
||||
for name, mod in module.named_children():
|
||||
child_name = f"{prefix}{name}"
|
||||
_partition_module(
|
||||
child_name,
|
||||
child_name + ".",
|
||||
module=mod,
|
||||
device_mesh=device_mesh,
|
||||
func=func,
|
||||
)
|
||||
|
||||
func(mod_name, module, device_mesh)
|
||||
|
||||
|
||||
def partition_module(
|
||||
module: nn.Module,
|
||||
device_mesh: DeviceMesh,
|
||||
func: Callable,
|
||||
to_local: bool = False,
|
||||
):
|
||||
"""partition module.
|
||||
|
||||
Parameters in module won't be force Replicated.
|
||||
|
||||
Args:
|
||||
module (Module): Module to be partitioned.
|
||||
device_mesh (DeviceMesh): The device mesh.
|
||||
func (Callable): partition callback.
|
||||
to_local (bool): Convert all DTensor parameters to Tensor parameters.
|
||||
"""
|
||||
_partition_module(
|
||||
"", "", module=module, device_mesh=device_mesh, func=func
|
||||
)
|
||||
|
||||
if to_local:
|
||||
module_to_local(module)
|
||||
|
||||
|
||||
def replicate_module(model: nn.Module, device_mesh: DeviceMesh):
|
||||
"""Replicate all parameters in module.
|
||||
|
||||
Args:
|
||||
model (Module): Module to perform replicate.
|
||||
device_mesh (DeviceMesh): The distribution device mesh.
|
||||
"""
|
||||
for name, param in model.named_parameters(recurse=False):
|
||||
param = distribute_tensor(
|
||||
param, device_mesh=device_mesh, placements=[Replicate()]
|
||||
).to_local()
|
||||
param = nn.Parameter(param)
|
||||
model.register_parameter(name, param)
|
||||
|
||||
for name, buf in model.named_buffers(recurse=False):
|
||||
buf = distribute_tensor(
|
||||
buf, device_mesh=device_mesh, placements=[Replicate()]
|
||||
).to_local()
|
||||
model.register_buffer(name, buf)
|
||||
@ -0,0 +1,79 @@
|
||||
import torch
|
||||
|
||||
|
||||
def continuous_tensor(
|
||||
inputs: torch.Tensor, seq_length: torch.LongTensor
|
||||
):
|
||||
"""Convert batched tensor to continuous tensor.
|
||||
|
||||
Args:
|
||||
inputs (Tensor): batched tensor.
|
||||
seq_length (Tensor): length of each sequence.
|
||||
|
||||
Return:
|
||||
Tensor: continuoused tensor.
|
||||
"""
|
||||
assert inputs.dim() > 1
|
||||
if inputs.size(1) == 1:
|
||||
return inputs.reshape(1, -1)
|
||||
|
||||
inputs = [inp[:slen] for inp, slen in zip(inputs, seq_length)]
|
||||
|
||||
inputs = torch.cat(inputs).unsqueeze(0)
|
||||
return inputs
|
||||
|
||||
|
||||
def batch_tensor(inputs: torch.Tensor, seq_length: torch.LongTensor):
|
||||
"""Convert continuoused tensor to batched tensor.
|
||||
|
||||
Args:
|
||||
inputs (Tensor): continuoused tensor.
|
||||
seq_length (Tensor): length of each sequence.
|
||||
|
||||
Return:
|
||||
Tensor: batched tensor.
|
||||
"""
|
||||
from torch.nn.utils.rnn import pad_sequence
|
||||
|
||||
end_loc = seq_length.cumsum(0)
|
||||
start_loc = end_loc - seq_length
|
||||
|
||||
inputs = [
|
||||
inputs[0, sloc:eloc] for sloc, eloc in zip(start_loc, end_loc)
|
||||
]
|
||||
inputs = pad_sequence(inputs, batch_first=True)
|
||||
return inputs
|
||||
|
||||
|
||||
def page_cache(
|
||||
paged_cache: torch.Tensor,
|
||||
batched_cache: torch.Tensor,
|
||||
cache_length: torch.Tensor,
|
||||
block_offsets: torch.Tensor,
|
||||
permute_head: bool = True,
|
||||
):
|
||||
"""Convert batched cache to paged cache.
|
||||
|
||||
Args:
|
||||
paged_cache (Tensor): Output paged cache.
|
||||
batched_cache (Tensor): Input batched cache.
|
||||
cache_length (Tensor): length of the cache.
|
||||
block_offsets (Tensor): Offset of each blocks.
|
||||
"""
|
||||
assert block_offsets.dim() == 2
|
||||
block_size = paged_cache.size(1)
|
||||
batch_size = batched_cache.size(0)
|
||||
if permute_head:
|
||||
batched_cache = batched_cache.permute(0, 2, 1, 3)
|
||||
|
||||
for b_idx in range(batch_size):
|
||||
cache_len = cache_length[b_idx]
|
||||
b_cache = batched_cache[b_idx]
|
||||
block_off = block_offsets[b_idx]
|
||||
block_off_idx = 0
|
||||
for s_start in range(0, cache_len, block_size):
|
||||
s_end = min(s_start + block_size, cache_len)
|
||||
s_len = s_end - s_start
|
||||
b_off = block_off[block_off_idx]
|
||||
paged_cache[b_off, :s_len] = b_cache[s_start:s_end]
|
||||
block_off_idx += 1
|
||||
@ -1,71 +0,0 @@
|
||||
import pytest
|
||||
from swarms.models.cog_agent import CogAgent
|
||||
from unittest.mock import MagicMock
|
||||
from PIL import Image
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def cogagent_params():
|
||||
return {
|
||||
"model_name": "ZhipuAI/cogagent-chat",
|
||||
"tokenizer_name": "I-ModelScope/vicuna-7b-v1.5",
|
||||
"dtype": "torch.bfloat16",
|
||||
"low_cpu_mem_usage": True,
|
||||
"load_in_4bit": True,
|
||||
"trust_remote_code": True,
|
||||
"device": "cuda",
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def cogagent(cogagent_params):
|
||||
return CogAgent(**cogagent_params)
|
||||
|
||||
|
||||
def test_init(mocker, cogagent_params, cogagent):
|
||||
mock_model = mocker.patch(
|
||||
"swarms.models.cog_agent.AutoModelForCausalLM.from_pretrained"
|
||||
)
|
||||
mock_tokenizer = mocker.patch(
|
||||
"swarms.models.cog_agent.AutoTokenizer.from_pretrained"
|
||||
)
|
||||
|
||||
for param, value in cogagent_params.items():
|
||||
assert getattr(cogagent, param) == value
|
||||
|
||||
mock_tokenizer.assert_called_once_with(
|
||||
cogagent_params["tokenizer_name"]
|
||||
)
|
||||
mock_model.assert_called_once_with(
|
||||
cogagent_params["model_name"],
|
||||
torch_dtype=cogagent_params["dtype"],
|
||||
low_cpu_mem_usage=cogagent_params["low_cpu_mem_usage"],
|
||||
load_in_4bit=cogagent_params["load_in_4bit"],
|
||||
trust_remote_code=cogagent_params["trust_remote_code"],
|
||||
)
|
||||
|
||||
|
||||
def test_run(mocker, cogagent):
|
||||
task = "How are you?"
|
||||
img = "images/1.jpg"
|
||||
mock_image = mocker.patch(
|
||||
"PIL.Image.open", return_value=MagicMock(spec=Image.Image)
|
||||
)
|
||||
cogagent.model.build_conversation_input_ids = MagicMock(
|
||||
return_value={
|
||||
"input_ids": MagicMock(),
|
||||
"token_type_ids": MagicMock(),
|
||||
"attention_mask": MagicMock(),
|
||||
"images": [MagicMock()],
|
||||
}
|
||||
)
|
||||
cogagent.model.__call__ = MagicMock(return_value="Mocked output")
|
||||
cogagent.decode = MagicMock(return_value="Mocked response")
|
||||
|
||||
output = cogagent.run(task, img)
|
||||
|
||||
assert output is not None
|
||||
mock_image.assert_called_once_with(img)
|
||||
cogagent.model.build_conversation_input_ids.assert_called_once()
|
||||
cogagent.model.__call__.assert_called_once()
|
||||
cogagent.decode.assert_called_once()
|
||||
@ -1,39 +0,0 @@
|
||||
import pytest
|
||||
from swarms.models.modelscope_pipeline import ModelScopePipeline
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def pipeline_params():
|
||||
return {
|
||||
"type_task": "text-generation",
|
||||
"model_name": "gpt2",
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def pipeline_model(pipeline_params):
|
||||
return ModelScopePipeline(**pipeline_params)
|
||||
|
||||
|
||||
def test_init(mocker, pipeline_params, pipeline_model):
|
||||
mock_pipeline = mocker.patch(
|
||||
"swarms.models.modelscope_pipeline.pipeline"
|
||||
)
|
||||
|
||||
for param, value in pipeline_params.items():
|
||||
assert getattr(pipeline_model, param) == value
|
||||
|
||||
mock_pipeline.assert_called_once_with(
|
||||
pipeline_params["type_task"],
|
||||
model=pipeline_params["model_name"],
|
||||
)
|
||||
|
||||
|
||||
def test_run(mocker, pipeline_model):
|
||||
task = "Generate a 10,000 word blog on health and wellness."
|
||||
pipeline_model.model = MagicMock(return_value="Mocked output")
|
||||
|
||||
output = pipeline_model.run(task)
|
||||
|
||||
assert output is not None
|
||||
@ -1,58 +0,0 @@
|
||||
import pytest
|
||||
from swarms.models.modelscope_llm import ModelScopeAutoModel
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def model_params():
|
||||
return {
|
||||
"model_name": "gpt2",
|
||||
"tokenizer_name": None,
|
||||
"device": "cuda",
|
||||
"device_map": "auto",
|
||||
"max_new_tokens": 500,
|
||||
"skip_special_tokens": True,
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def modelscope(model_params):
|
||||
return ModelScopeAutoModel(**model_params)
|
||||
|
||||
|
||||
def test_init(mocker, model_params, modelscope):
|
||||
mock_model = mocker.patch(
|
||||
"swarms.models.modelscope_llm.AutoModelForCausalLM.from_pretrained"
|
||||
)
|
||||
mock_tokenizer = mocker.patch(
|
||||
"swarms.models.modelscope_llm.AutoTokenizer.from_pretrained"
|
||||
)
|
||||
|
||||
for param, value in model_params.items():
|
||||
assert getattr(modelscope, param) == value
|
||||
|
||||
mock_tokenizer.assert_called_once_with(
|
||||
model_params["tokenizer_name"]
|
||||
)
|
||||
mock_model.assert_called_once_with(
|
||||
model_params["model_name"],
|
||||
device_map=model_params["device_map"],
|
||||
)
|
||||
|
||||
|
||||
def test_run(mocker, modelscope):
|
||||
task = "Generate a 10,000 word blog on health and wellness."
|
||||
mocker.patch(
|
||||
"swarms.models.modelscope_llm.AutoTokenizer.decode",
|
||||
return_value="Mocked output",
|
||||
)
|
||||
modelscope.model.generate = MagicMock(
|
||||
return_value=["Mocked token"]
|
||||
)
|
||||
modelscope.tokenizer = MagicMock(
|
||||
return_value={"input_ids": "Mocked input_ids"}
|
||||
)
|
||||
|
||||
output = modelscope.run(task)
|
||||
|
||||
assert output is not None
|
||||
@ -1,141 +0,0 @@
|
||||
import pytest
|
||||
from swarms.models.vllm import vLLM
|
||||
|
||||
|
||||
# Fixture for initializing vLLM
|
||||
@pytest.fixture
|
||||
def vllm_instance():
|
||||
return vLLM()
|
||||
|
||||
|
||||
# Test the default initialization of vLLM
|
||||
def test_vllm_default_init(vllm_instance):
|
||||
assert isinstance(vllm_instance, vLLM)
|
||||
assert vllm_instance.model_name == "facebook/opt-13b"
|
||||
assert vllm_instance.tensor_parallel_size == 4
|
||||
assert not vllm_instance.trust_remote_code
|
||||
assert vllm_instance.revision is None
|
||||
assert vllm_instance.temperature == 0.5
|
||||
assert vllm_instance.top_p == 0.95
|
||||
|
||||
|
||||
# Test custom initialization of vLLM
|
||||
def test_vllm_custom_init():
|
||||
vllm_instance = vLLM(
|
||||
model_name="custom_model",
|
||||
tensor_parallel_size=8,
|
||||
trust_remote_code=True,
|
||||
revision="123",
|
||||
temperature=0.7,
|
||||
top_p=0.9,
|
||||
)
|
||||
assert isinstance(vllm_instance, vLLM)
|
||||
assert vllm_instance.model_name == "custom_model"
|
||||
assert vllm_instance.tensor_parallel_size == 8
|
||||
assert vllm_instance.trust_remote_code
|
||||
assert vllm_instance.revision == "123"
|
||||
assert vllm_instance.temperature == 0.7
|
||||
assert vllm_instance.top_p == 0.9
|
||||
|
||||
|
||||
# Test the run method of vLLM
|
||||
def test_vllm_run(vllm_instance):
|
||||
task = "Hello, vLLM!"
|
||||
result = vllm_instance.run(task)
|
||||
assert isinstance(result, str)
|
||||
assert len(result) > 0
|
||||
|
||||
|
||||
# Test run method with different temperature and top_p values
|
||||
@pytest.mark.parametrize(
|
||||
"temperature, top_p", [(0.2, 0.8), (0.8, 0.2)]
|
||||
)
|
||||
def test_vllm_run_with_params(vllm_instance, temperature, top_p):
|
||||
task = "Temperature and Top-P Test"
|
||||
result = vllm_instance.run(
|
||||
task, temperature=temperature, top_p=top_p
|
||||
)
|
||||
assert isinstance(result, str)
|
||||
assert len(result) > 0
|
||||
|
||||
|
||||
# Test run method with a specific model revision
|
||||
def test_vllm_run_with_revision(vllm_instance):
|
||||
task = "Specific Model Revision Test"
|
||||
result = vllm_instance.run(task, revision="abc123")
|
||||
assert isinstance(result, str)
|
||||
assert len(result) > 0
|
||||
|
||||
|
||||
# Test run method with a specific model name
|
||||
def test_vllm_run_with_custom_model(vllm_instance):
|
||||
task = "Custom Model Test"
|
||||
custom_model_name = "my_custom_model"
|
||||
result = vllm_instance.run(task, model_name=custom_model_name)
|
||||
assert isinstance(result, str)
|
||||
assert len(result) > 0
|
||||
assert vllm_instance.model_name == custom_model_name
|
||||
|
||||
|
||||
# Test run method with invalid task input
|
||||
def test_vllm_run_invalid_task(vllm_instance):
|
||||
invalid_task = None
|
||||
with pytest.raises(ValueError):
|
||||
vllm_instance.run(invalid_task)
|
||||
|
||||
|
||||
# Test run method with a very high temperature value
|
||||
def test_vllm_run_high_temperature(vllm_instance):
|
||||
task = "High Temperature Test"
|
||||
high_temperature = 10.0
|
||||
result = vllm_instance.run(task, temperature=high_temperature)
|
||||
assert isinstance(result, str)
|
||||
assert len(result) > 0
|
||||
|
||||
|
||||
# Test run method with a very low top_p value
|
||||
def test_vllm_run_low_top_p(vllm_instance):
|
||||
task = "Low Top-P Test"
|
||||
low_top_p = 0.01
|
||||
result = vllm_instance.run(task, top_p=low_top_p)
|
||||
assert isinstance(result, str)
|
||||
assert len(result) > 0
|
||||
|
||||
|
||||
# Test run method with an empty task
|
||||
def test_vllm_run_empty_task(vllm_instance):
|
||||
empty_task = ""
|
||||
result = vllm_instance.run(empty_task)
|
||||
assert isinstance(result, str)
|
||||
assert len(result) == 0
|
||||
|
||||
|
||||
# Test initialization with invalid parameters
|
||||
def test_vllm_invalid_init():
|
||||
with pytest.raises(ValueError):
|
||||
vLLM(
|
||||
model_name=None,
|
||||
tensor_parallel_size=-1,
|
||||
trust_remote_code="invalid",
|
||||
revision=123,
|
||||
temperature=-0.1,
|
||||
top_p=1.1,
|
||||
)
|
||||
|
||||
|
||||
# Test running vLLM with a large number of parallel heads
|
||||
def test_vllm_large_parallel_heads():
|
||||
vllm_instance = vLLM(tensor_parallel_size=16)
|
||||
task = "Large Parallel Heads Test"
|
||||
result = vllm_instance.run(task)
|
||||
assert isinstance(result, str)
|
||||
assert len(result) > 0
|
||||
|
||||
|
||||
# Test running vLLM with trust_remote_code set to True
|
||||
def test_vllm_trust_remote_code():
|
||||
vllm_instance = vLLM(trust_remote_code=True)
|
||||
task = "Trust Remote Code Test"
|
||||
result = vllm_instance.run(task)
|
||||
assert isinstance(result, str)
|
||||
assert len(result) > 0
|
||||
Loading…
Reference in new issue