724 lines
31 KiB
Python
724 lines
31 KiB
Python
import json
|
|
import os
|
|
import copy
|
|
from dataclasses import dataclass
|
|
from typing import Any, Optional, Tuple, Union
|
|
import torch
|
|
from transformers.configuration_utils import PretrainedConfig
|
|
from transformers.utils import is_torch_available
|
|
from transformers.cache_utils import QuantizedCache, StaticCache, DynamicCache
|
|
from transformers.generation.logits_process import LogitsProcessor
|
|
from transformers.pytorch_utils import Conv1D, prune_linear_layer
|
|
from torch import nn
|
|
from math import ceil
|
|
import tempfile
|
|
import huggingface_hub
|
|
from huggingface_hub.file_download import http_get
|
|
from urllib.parse import urlparse
|
|
|
|
FLAX_WEIGHTS_NAME = "flax_model.msgpack"
|
|
TF2_WEIGHTS_NAME = "tf_model.h5"
|
|
TF_WEIGHTS_NAME = "model.ckpt"
|
|
|
|
|
|
def assert_device_map(device_map, num_blocks):
|
|
blocks = list(range(0, num_blocks))
|
|
|
|
device_map_blocks = [item for sublist in list(device_map.values()) for item in sublist]
|
|
|
|
# Duplicate check
|
|
duplicate_blocks = []
|
|
for i in device_map_blocks:
|
|
if device_map_blocks.count(i) > 1 and i not in duplicate_blocks:
|
|
duplicate_blocks.append(i)
|
|
# Missing blocks
|
|
missing_blocks = [i for i in blocks if i not in device_map_blocks]
|
|
extra_blocks = [i for i in device_map_blocks if i not in blocks]
|
|
|
|
if len(duplicate_blocks) != 0:
|
|
raise ValueError(
|
|
"Duplicate attention blocks specified in device_map. Attention blocks must be specified to one device."
|
|
" These attention blocks were specified more than once: " + str(duplicate_blocks)
|
|
)
|
|
if len(missing_blocks) != 0:
|
|
raise ValueError(
|
|
"There are attention blocks for this model that are not specified in the device_map. Add these attention "
|
|
"blocks to a device on the device_map: " + str(missing_blocks)
|
|
)
|
|
if len(extra_blocks) != 0:
|
|
raise ValueError(
|
|
"The device_map contains more attention blocks than this model has. Remove these from the device_map:"
|
|
+ str(extra_blocks)
|
|
)
|
|
|
|
|
|
def get_device_map(n_layers, devices):
|
|
"""Returns a dictionary of layers distributed evenly across all devices."""
|
|
layers = list(range(n_layers))
|
|
n_blocks = int(ceil(n_layers / len(devices)))
|
|
layers_list = [layers[i : i + n_blocks] for i in range(0, n_layers, n_blocks)]
|
|
|
|
return dict(zip(devices, layers_list))
|
|
|
|
def isin_mps_friendly(elements: torch.Tensor, test_elements: torch.Tensor | int) -> torch.Tensor:
|
|
"""
|
|
Same as `torch.isin` without flags, but MPS-friendly. We can remove this function when we stop supporting
|
|
torch <= 2.3. See https://github.com/pytorch/pytorch/issues/77764#issuecomment-2067838075
|
|
|
|
Args:
|
|
elements (`torch.Tensor`): Input elements
|
|
test_elements (`torch.Tensor` or `int`): The elements to check against.
|
|
|
|
Returns:
|
|
`torch.Tensor`: A boolean tensor of the same shape as `elements` that is True for `elements` in `test_elements`
|
|
and False otherwise
|
|
"""
|
|
|
|
if elements.device.type == "mps" and not is_torch_greater_or_equal_than_2_4:
|
|
test_elements = torch.tensor(test_elements)
|
|
if test_elements.ndim == 0:
|
|
test_elements = test_elements.unsqueeze(0)
|
|
return elements.tile(test_elements.shape[0], 1).eq(test_elements.unsqueeze(1)).sum(dim=0).bool().squeeze()
|
|
else:
|
|
# Note: don't use named arguments in `torch.isin`, see https://github.com/pytorch/pytorch/issues/126045
|
|
return torch.isin(elements, test_elements)
|
|
|
|
def find_pruneable_heads_and_indices(
|
|
heads: list[int], n_heads: int, head_size: int, already_pruned_heads: set[int]
|
|
) -> tuple[set[int], torch.LongTensor]:
|
|
"""
|
|
Finds the heads and their indices taking `already_pruned_heads` into account.
|
|
|
|
Args:
|
|
heads (`list[int]`): List of the indices of heads to prune.
|
|
n_heads (`int`): The number of heads in the model.
|
|
head_size (`int`): The size of each head.
|
|
already_pruned_heads (`Set[int]`): A set of already pruned heads.
|
|
|
|
Returns:
|
|
`tuple[Set[int], torch.LongTensor]`: A tuple with the indices of heads to prune taking `already_pruned_heads`
|
|
into account and the indices of rows/columns to keep in the layer weight.
|
|
"""
|
|
mask = torch.ones(n_heads, head_size)
|
|
heads = set(heads) - already_pruned_heads # Convert to set and remove already pruned heads
|
|
for head in heads:
|
|
# Compute how many pruned heads are before the head and move the index accordingly
|
|
head = head - sum(1 if h < head else 0 for h in already_pruned_heads)
|
|
mask[head] = 0
|
|
mask = mask.view(-1).contiguous().eq(1)
|
|
index: torch.LongTensor = torch.arange(len(mask))[mask].long()
|
|
return heads, index
|
|
|
|
def prune_conv1d_layer(layer: Conv1D, index: torch.LongTensor, dim: int = 1) -> Conv1D:
|
|
"""
|
|
Prune a Conv1D layer to keep only entries in index. A Conv1D work as a Linear layer (see e.g. BERT) but the weights
|
|
are transposed.
|
|
|
|
Used to remove heads.
|
|
|
|
Args:
|
|
layer ([`~pytorch_utils.Conv1D`]): The layer to prune.
|
|
index (`torch.LongTensor`): The indices to keep in the layer.
|
|
dim (`int`, *optional*, defaults to 1): The dimension on which to keep the indices.
|
|
|
|
Returns:
|
|
[`~pytorch_utils.Conv1D`]: The pruned layer as a new layer with `requires_grad=True`.
|
|
"""
|
|
index = index.to(layer.weight.device)
|
|
W = layer.weight.index_select(dim, index).detach().clone()
|
|
if dim == 0:
|
|
b = layer.bias.detach().clone()
|
|
else:
|
|
b = layer.bias[index].detach().clone()
|
|
new_size = list(layer.weight.size())
|
|
new_size[dim] = len(index)
|
|
new_layer = Conv1D(new_size[1], new_size[0]).to(layer.weight.device)
|
|
new_layer.weight.requires_grad = False
|
|
new_layer.weight.copy_(W.contiguous())
|
|
new_layer.weight.requires_grad = True
|
|
new_layer.bias.requires_grad = False
|
|
new_layer.bias.copy_(b.contiguous())
|
|
new_layer.bias.requires_grad = True
|
|
return new_layer
|
|
|
|
def download_url(url, proxies=None):
|
|
"""
|
|
Downloads a given url in a temporary file. This function is not safe to use in multiple processes. Its only use is
|
|
for deprecated behavior allowing to download config/models with a single url instead of using the Hub.
|
|
|
|
Args:
|
|
url (`str`): The url of the file to download.
|
|
proxies (`dict[str, str]`, *optional*):
|
|
A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
|
|
'http://hostname': 'foo.bar:4012'}.` The proxies are used on each request.
|
|
|
|
Returns:
|
|
`str`: The location of the temporary file where the url was downloaded.
|
|
"""
|
|
tmp_fd, tmp_file = tempfile.mkstemp()
|
|
with os.fdopen(tmp_fd, "wb") as f:
|
|
http_get(url, f, proxies=proxies)
|
|
return tmp_file
|
|
|
|
_is_offline_mode = huggingface_hub.constants.HF_HUB_OFFLINE
|
|
|
|
def is_offline_mode():
|
|
return _is_offline_mode
|
|
|
|
def is_remote_url(url_or_filename):
|
|
parsed = urlparse(url_or_filename)
|
|
return parsed.scheme in ("http", "https")
|
|
|
|
def is_safetensors_available() -> Union[tuple[bool, str], bool]:
|
|
return True
|
|
|
|
def is_torch_sdpa_available():
|
|
# Mostly retained for backward compatibility in remote code, since sdpa works correctly on all torch versions >= 2.2
|
|
return True
|
|
|
|
def prune_layer(layer: nn.Linear | Conv1D, index: torch.LongTensor, dim: int | None = None) -> nn.Linear | Conv1D:
|
|
"""
|
|
Prune a Conv1D or linear layer to keep only entries in index.
|
|
|
|
Used to remove heads.
|
|
|
|
Args:
|
|
layer (`Union[torch.nn.Linear, Conv1D]`): The layer to prune.
|
|
index (`torch.LongTensor`): The indices to keep in the layer.
|
|
dim (`int`, *optional*): The dimension on which to keep the indices.
|
|
|
|
Returns:
|
|
`torch.nn.Linear` or [`~pytorch_utils.Conv1D`]: The pruned layer as a new layer with `requires_grad=True`.
|
|
"""
|
|
if isinstance(layer, nn.Linear):
|
|
return prune_linear_layer(layer, index, dim=0 if dim is None else dim)
|
|
elif isinstance(layer, Conv1D):
|
|
return prune_conv1d_layer(layer, index, dim=1 if dim is None else dim)
|
|
else:
|
|
raise ValueError(f"Can't prune layer of class {layer.__class__}")
|
|
|
|
class HammingDiversityLogitsProcessor(LogitsProcessor):
|
|
r"""
|
|
[`LogitsProcessor`] that enforces diverse beam search.
|
|
Note that this logits processor is only effective for [`PreTrainedModel.group_beam_search`]. See [Diverse Beam
|
|
Search: Decoding Diverse Solutions from Neural Sequence Models](https://huggingface.co/papers/1610.02424) for more
|
|
details.
|
|
Traditional beam search often generates very similar sequences across different beams.
|
|
`HammingDiversityLogitsProcessor` addresses this by penalizing beams that generate tokens already chosen by other
|
|
beams in the same time step.
|
|
Args:
|
|
diversity_penalty (`float`):
|
|
This value is subtracted from a beam's score if it generates a token same as any beam from other group at a
|
|
particular time. A higher `diversity_penalty` will enforce greater diversity among the beams. Adjusting
|
|
this value can help strike a balance between diversity and natural likelihood.
|
|
num_beams (`int`):
|
|
Number of beams for beam search. 1 means no beam search.
|
|
num_beam_groups (`int`):
|
|
Number of groups to divide `num_beams` into in order to ensure diversity among different groups of beams.
|
|
[this paper](https://huggingface.co/papers/1610.02424) for more details.
|
|
Examples:
|
|
```python
|
|
>>> from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
|
|
>>> import torch
|
|
>>> # Initialize the model and tokenizer
|
|
>>> tokenizer = AutoTokenizer.from_pretrained("google-t5/t5-base")
|
|
>>> model = AutoModelForSeq2SeqLM.from_pretrained("google-t5/t5-base")
|
|
>>> # A long text about the solar system
|
|
>>> text = (
|
|
... "The Solar System is a gravitationally bound system comprising the Sun and the objects that orbit it, "
|
|
... "either directly or indirectly. Of the objects that orbit the Sun directly, the largest are the eight "
|
|
... "planets, with the remainder being smaller objects, such as the five dwarf planets and small Solar System "
|
|
... "bodies. The Solar System formed 4.6 billion years ago from the gravitational collapse of a giant "
|
|
... "interstellar molecular cloud."
|
|
... )
|
|
>>> inputs = tokenizer("summarize: " + text, return_tensors="pt")
|
|
>>> # Generate diverse summary
|
|
>>> outputs_diverse = model.generate(
|
|
... **inputs,
|
|
... num_beam_groups=2,
|
|
... diversity_penalty=10.0,
|
|
... max_length=100,
|
|
... num_beams=4,
|
|
... num_return_sequences=2,
|
|
... )
|
|
>>> summaries_diverse = tokenizer.batch_decode(outputs_diverse, skip_special_tokens=True)
|
|
>>> # Generate non-diverse summary
|
|
>>> outputs_non_diverse = model.generate(
|
|
... **inputs,
|
|
... max_length=100,
|
|
... num_beams=4,
|
|
... num_return_sequences=2,
|
|
... )
|
|
>>> summary_non_diverse = tokenizer.batch_decode(outputs_non_diverse, skip_special_tokens=True)
|
|
>>> # With `diversity_penalty`, the resulting beams are much more diverse
|
|
>>> print(summary_non_diverse)
|
|
['the solar system formed 4.6 billion years ago from the collapse of a giant interstellar molecular cloud. of the objects that orbit the Sun directly, the largest are the eight planets.',
|
|
'the Solar System formed 4.6 billion years ago from the collapse of a giant interstellar molecular cloud. of the objects that orbit the Sun directly, the largest are the eight planets.']
|
|
>>> print(summaries_diverse)
|
|
['the solar system formed 4.6 billion years ago from the collapse of a giant interstellar molecular cloud. of the objects that orbit the Sun directly, the largest are the eight planets.',
|
|
'the solar system formed 4.6 billion years ago from the collapse of a giant interstellar molecular cloud. of the objects that orbit the Sun directly, the largest are the eight planets. the rest of the objects are smaller objects, such as the five dwarf planets and small solar system bodies.']
|
|
```
|
|
"""
|
|
|
|
def __init__(self, diversity_penalty: float, num_beams: int, num_beam_groups: int):
|
|
if not isinstance(diversity_penalty, float) or (not diversity_penalty > 0.0):
|
|
raise ValueError("`diversity_penalty` should be a float strictly larger than 0.")
|
|
self._diversity_penalty = diversity_penalty
|
|
if not isinstance(num_beams, int) or num_beams < 2:
|
|
raise ValueError("`num_beams` should be an integer strictly larger than 1.")
|
|
self._num_beams = num_beams
|
|
if not isinstance(num_beam_groups, int) or num_beam_groups < 2:
|
|
raise ValueError("`num_beam_groups` should be an integer strictly larger than 1.")
|
|
if num_beam_groups > num_beams:
|
|
raise ValueError("`beam_groups` has to be smaller or equal to `num_beams`.")
|
|
self._num_sub_beams = num_beams // num_beam_groups
|
|
|
|
def __call__(
|
|
self,
|
|
input_ids: torch.LongTensor,
|
|
scores: torch.FloatTensor,
|
|
current_tokens: torch.LongTensor,
|
|
beam_group_idx: int,
|
|
) -> torch.FloatTensor:
|
|
r"""
|
|
Args:
|
|
input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
|
|
Indices of input sequence tokens in the vocabulary. [What are input IDs?](../glossary#input-ids)
|
|
scores (`torch.FloatTensor` of shape `(batch_size, config.vocab_size)`):
|
|
Prediction scores of a language modeling head. These can be logits for each vocabulary when not using
|
|
beam search or log softmax for each vocabulary token when using beam search
|
|
current_tokens (`torch.LongTensor` of shape `(batch_size)`):
|
|
Indices of input sequence tokens in the vocabulary, corresponding to the tokens selected by the other
|
|
beam groups in the current generation step.
|
|
beam_group_idx (`int`):
|
|
The index of the beam group currently being processed.
|
|
Return:
|
|
`torch.FloatTensor` of shape `(batch_size, config.vocab_size)`:
|
|
The processed prediction scores.
|
|
"""
|
|
# hamming diversity: penalise using same token in current group which was used in previous groups at
|
|
# the same time step
|
|
batch_size = current_tokens.shape[0] // self._num_beams
|
|
group_start_idx = beam_group_idx * self._num_sub_beams
|
|
group_end_idx = min(group_start_idx + self._num_sub_beams, self._num_beams)
|
|
group_size = group_end_idx - group_start_idx
|
|
vocab_size = scores.shape[-1]
|
|
|
|
if group_start_idx == 0:
|
|
return scores
|
|
|
|
scores_processed = scores.clone()
|
|
for batch_idx in range(batch_size):
|
|
# predicted tokens of last time step of previous groups
|
|
previous_group_tokens = current_tokens[
|
|
batch_idx * self._num_beams : batch_idx * self._num_beams + group_start_idx
|
|
]
|
|
token_frequency = torch.bincount(previous_group_tokens, minlength=vocab_size).to(scores.device)
|
|
scores_processed[batch_idx * group_size : (batch_idx + 1) * group_size] -= (
|
|
self._diversity_penalty * token_frequency
|
|
)
|
|
|
|
return scores_processed
|
|
|
|
@dataclass
|
|
class CacheConfig:
|
|
"""
|
|
Base class for cache configs
|
|
"""
|
|
|
|
cache_implementation: None
|
|
|
|
@classmethod
|
|
def from_dict(cls, config_dict, **kwargs):
|
|
"""
|
|
Constructs a CacheConfig instance from a dictionary of parameters.
|
|
Args:
|
|
config_dict (dict[str, Any]): Dictionary containing configuration parameters.
|
|
**kwargs: Additional keyword arguments to override dictionary values.
|
|
|
|
Returns:
|
|
CacheConfig: Instance of CacheConfig constructed from the dictionary.
|
|
"""
|
|
config = cls(**config_dict)
|
|
to_remove = []
|
|
for key, value in kwargs.items():
|
|
if hasattr(config, key):
|
|
setattr(config, key, value)
|
|
to_remove.append(key)
|
|
for key in to_remove:
|
|
kwargs.pop(key, None)
|
|
return config
|
|
|
|
# Copied from transformers.utils.quantization_config.QuantizationConfigMixin.to_json_file
|
|
def to_json_file(self, json_file_path: Union[str, os.PathLike]):
|
|
"""
|
|
Save this instance to a JSON file.
|
|
|
|
Args:
|
|
json_file_path (`str` or `os.PathLike`):
|
|
Path to the JSON file in which this configuration instance's parameters will be saved.
|
|
use_diff (`bool`, *optional*, defaults to `True`):
|
|
If set to `True`, only the difference between the config instance and the default
|
|
`QuantizationConfig()` is serialized to JSON file.
|
|
"""
|
|
with open(json_file_path, "w", encoding="utf-8") as writer:
|
|
config_dict = self.to_dict()
|
|
json_string = json.dumps(config_dict, indent=2, sort_keys=True) + "\n"
|
|
|
|
writer.write(json_string)
|
|
|
|
# Copied from transformers.utils.quantization_config.QuantizationConfigMixin.to_dict
|
|
def to_dict(self) -> dict[str, Any]:
|
|
"""
|
|
Serializes this instance to a Python dictionary. Returns:
|
|
`dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance.
|
|
"""
|
|
return copy.deepcopy(self.__dict__)
|
|
|
|
# Copied from transformers.utils.quantization_config.QuantizationConfigMixin.__iter__
|
|
def __iter__(self):
|
|
"""allows `dict(obj)` for situations where obj may be a dict or QuantizationConfigMixin"""
|
|
for attr, value in copy.deepcopy(self.__dict__).items():
|
|
yield attr, value
|
|
|
|
# Copied from transformers.utils.quantization_config.QuantizationConfigMixin.__repr__
|
|
def __repr__(self):
|
|
return f"{self.__class__.__name__} {self.to_json_string()}"
|
|
|
|
def to_json_string(self):
|
|
"""
|
|
Serializes this instance to a JSON formatted string.
|
|
Returns:
|
|
str: JSON formatted string representing the configuration instance.
|
|
"""
|
|
return json.dumps(self.__dict__, indent=2) + "\n"
|
|
|
|
# Copied from transformers.utils.quantization_config.QuantizationConfigMixin.update
|
|
def update(self, **kwargs):
|
|
"""
|
|
Updates attributes of this class instance with attributes from `kwargs` if they match existing attributes,
|
|
returning all the unused kwargs.
|
|
|
|
Args:
|
|
kwargs (`dict[str, Any]`):
|
|
Dictionary of attributes to tentatively update this class.
|
|
|
|
Returns:
|
|
`dict[str, Any]`: Dictionary containing all the key-value pairs that were not used to update the instance.
|
|
"""
|
|
to_remove = []
|
|
for key, value in kwargs.items():
|
|
if hasattr(self, key):
|
|
setattr(self, key, value)
|
|
to_remove.append(key)
|
|
|
|
# Remove all the attributes that were updated, without modifying the input dict
|
|
unused_kwargs = {key: value for key, value in kwargs.items() if key not in to_remove}
|
|
return unused_kwargs
|
|
|
|
@dataclass
|
|
class QuantizedCacheConfig(CacheConfig):
|
|
"""
|
|
Configuration class for quantized cache settings.
|
|
|
|
Attributes:
|
|
backend (`str`, *optional*, defaults to `"quanto"`):
|
|
Backend to use when performing quantization, Can be one of [`quanto`, `HQQ`]
|
|
nbits (`Optional[int]`, *optional*, defaults to 4):
|
|
Number of bits, can be 2 or 4 for the `quanto` backend and one of [1, 2, 3, 4, 8] for the `HQQ` backend. Defaults to 2.
|
|
axis_key (`int`, *optional*, defaults to 0):
|
|
Axis over which to perform grouping for the key tensors. Can be [0, -1] for `quanto` backend and [0, 1] for `HQQ` backend.
|
|
axis_value (`int`, *optional*, defaults to 0):
|
|
Axis over which to perform grouping for the value tensors. Can be [0, -1] for `quanto` backend and [0, 1] for `HQQ` backend.
|
|
q_group_size (`Optional[int]`, *optional*, defaults to 64):
|
|
Size of the quantization group, should be a divisor of the model's hidden dimension.
|
|
Defaults to 64.
|
|
residual_length (`Optional[int]`, *optional*, defaults to 128):
|
|
Length of the residual cache which will always be stored in original precision.
|
|
Defaults to 128.
|
|
compute_dtype (`torch.dtype`, *optional*, defaults to `torch.float16`):
|
|
The default dtype used for computations in the model. Keys and Values will be cast to this dtype after dequantization.
|
|
device (`str`, *optional*, defaults to `"cpu"`):
|
|
Device on which to perform computations, should be same as the model's device.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
backend: str = "quanto",
|
|
nbits: Optional[int] = 4,
|
|
axis_key: Optional[int] = 0,
|
|
axis_value: Optional[int] = 0,
|
|
q_group_size: Optional[int] = 64,
|
|
residual_length: Optional[int] = 128,
|
|
compute_dtype: Optional[torch.dtype] = torch.float16,
|
|
device: Optional[str] = "cpu",
|
|
):
|
|
self.backend = backend
|
|
self.nbits = nbits
|
|
self.axis_key = axis_key
|
|
self.axis_value = axis_value
|
|
self.q_group_size = q_group_size
|
|
self.residual_length = residual_length
|
|
self.compute_dtype = compute_dtype
|
|
self.device = device
|
|
|
|
def validate(self):
|
|
"""Validates if the arguments passed are correct"""
|
|
|
|
incorrect_arg_msg = (
|
|
"Some of the keys in `cache_config` are defined incorrectly. `{key}` should be {correct_value}` "
|
|
"but found {found_value}"
|
|
)
|
|
# Check that the values are reasonable in general (nbits, axis)
|
|
# Later in QuantizedCache init we check if they are supported for that particular backend
|
|
if self.nbits not in [1, 2, 3, 4, 8]:
|
|
raise ValueError(
|
|
incorrect_arg_msg.format(
|
|
key="nbits",
|
|
correct_value="2 or 4 or 8",
|
|
found_value=self.nbits,
|
|
),
|
|
)
|
|
if self.q_group_size <= 0:
|
|
raise ValueError(
|
|
incorrect_arg_msg.format(
|
|
key="q_group_size",
|
|
correct_value="a positive integer",
|
|
found_value=self.q_group_size,
|
|
),
|
|
)
|
|
if self.residual_length < 0:
|
|
raise ValueError(
|
|
incorrect_arg_msg.format(
|
|
key="residual_length",
|
|
correct_value="a positive integer",
|
|
found_value=self.residual_length,
|
|
),
|
|
)
|
|
|
|
if self.axis_key not in [0, 1, -1]:
|
|
raise ValueError(
|
|
incorrect_arg_msg.format(
|
|
key="axis_key",
|
|
correct_value="`1` or `0`, `-1`",
|
|
found_value=self.axis_key,
|
|
),
|
|
)
|
|
|
|
if self.axis_value not in [0, 1, -1]:
|
|
raise ValueError(
|
|
incorrect_arg_msg.format(
|
|
key="axis_value",
|
|
correct_value="`1` or `0` or `-1`",
|
|
found_value=self.axis_value,
|
|
),
|
|
)
|
|
|
|
@dataclass
|
|
class StaticCacheConfig(CacheConfig):
|
|
"""
|
|
Configuration class for static cache settings.
|
|
"""
|
|
|
|
cache_implementation = "static"
|
|
|
|
def __init__(self, batch_size: int, max_cache_len: int, device="cpu"):
|
|
self.batch_size = batch_size
|
|
self.max_cache_len = max_cache_len
|
|
self.device = device
|
|
|
|
def validate(self):
|
|
"""Validates if the arguments passed are correct"""
|
|
|
|
incorrect_arg_msg = (
|
|
"Some of the keys in `cache_config` are defined incorrectly. `{key}` should be {correct_value}` "
|
|
"but found {found_value}"
|
|
)
|
|
|
|
if self.batch_size <= 0:
|
|
raise ValueError(
|
|
incorrect_arg_msg.format(
|
|
key="batch_size",
|
|
correct_value="> 0",
|
|
found_value=self.batch_size,
|
|
),
|
|
)
|
|
|
|
if self.max_cache_len <= 0:
|
|
raise ValueError(
|
|
incorrect_arg_msg.format(
|
|
key="max_cache_len",
|
|
correct_value="> 0",
|
|
found_value=self.max_cache_len,
|
|
),
|
|
)
|
|
|
|
class MambaCache:
|
|
"""
|
|
Cache for mamba model which does not have attention mechanism and key value states.
|
|
|
|
Arguments:
|
|
config (`PretrainedConfig):
|
|
The configuration file defining the shape-related attributes required to initialize the static cache.
|
|
max_batch_size (`int`):
|
|
The maximum batch size with which the model will be used. Note that a new instance must be instantiated if a smaller batch size is used.
|
|
dtype (`torch.dtype`, *optional*, defaults to `torch.float16`):
|
|
The default `dtype` to use when initializing the layer.
|
|
device (`torch.device` or `str`, *optional*):
|
|
The device on which the cache should be initialized. Should be the same as the layer.
|
|
|
|
Example:
|
|
|
|
```python
|
|
>>> from transformers import AutoTokenizer, MambaForCausalLM, MambaCache
|
|
|
|
>>> model = MambaForCausalLM.from_pretrained("state-spaces/mamba-130m-hf")
|
|
>>> tokenizer = AutoTokenizer.from_pretrained("state-spaces/mamba-130m-hf")
|
|
|
|
>>> inputs = tokenizer(text="My name is Mamba", return_tensors="pt")
|
|
|
|
>>> # Prepare a cache class and pass it to model's forward
|
|
>>> past_key_values = MambaCache(config=model.config, max_batch_size=1, device=model.device, dtype=model.dtype)
|
|
>>> outputs = model(**inputs, past_key_values=past_key_values, use_cache=True)
|
|
>>> outputs.past_key_values
|
|
MambaCache()
|
|
```
|
|
"""
|
|
|
|
is_compileable = True
|
|
|
|
# TODO (joao): add layer_device_map arg and update code in `generate` accordingly
|
|
def __init__(
|
|
self,
|
|
config: PretrainedConfig,
|
|
max_batch_size: int,
|
|
dtype: torch.dtype = torch.float16,
|
|
device: Union[torch.device, str, None] = None,
|
|
):
|
|
self.max_batch_size = max_batch_size
|
|
self._dtype = dtype
|
|
self.intermediate_size = config.intermediate_size
|
|
self.ssm_state_size = config.state_size
|
|
self.conv_kernel_size = config.conv_kernel
|
|
|
|
self.conv_states: list[torch.Tensor] = []
|
|
self.ssm_states: list[torch.Tensor] = []
|
|
device = torch.device(device) if device is not None else None
|
|
for _ in range(config.num_hidden_layers):
|
|
conv_state: torch.Tensor = torch.zeros(
|
|
self.max_batch_size,
|
|
self.intermediate_size,
|
|
self.conv_kernel_size,
|
|
device=device,
|
|
dtype=self._dtype,
|
|
)
|
|
ssm_state: torch.Tensor = torch.zeros(
|
|
self.max_batch_size,
|
|
self.intermediate_size,
|
|
self.ssm_state_size,
|
|
device=device,
|
|
dtype=self._dtype,
|
|
)
|
|
|
|
torch._dynamo.mark_static_address(conv_state)
|
|
torch._dynamo.mark_static_address(ssm_state)
|
|
self.conv_states.append(conv_state)
|
|
self.ssm_states.append(ssm_state)
|
|
|
|
def update_conv_state(
|
|
self, layer_idx: int, new_conv_state: torch.Tensor, cache_position: torch.LongTensor
|
|
) -> torch.Tensor:
|
|
# This `if` blocks is only reached in multigpu and if `layer_device_map` is not passed. It is used
|
|
# when the cache is initialized in the forward pass (e.g. Mamba)
|
|
if self.conv_states[layer_idx].device != new_conv_state.device:
|
|
self.conv_states[layer_idx] = self.conv_states[layer_idx].to(new_conv_state.device)
|
|
|
|
conv_state = self.conv_states[layer_idx]
|
|
cache_position = cache_position.clamp(0, self.conv_kernel_size - 1)
|
|
|
|
conv_state = conv_state.roll(shifts=-1, dims=-1)
|
|
conv_state[:, :, cache_position] = new_conv_state.to(device=conv_state.device, dtype=conv_state.dtype)
|
|
self.conv_states[layer_idx].zero_()
|
|
self.conv_states[layer_idx] += conv_state
|
|
return self.conv_states[layer_idx]
|
|
|
|
def update_ssm_state(self, layer_idx: int, new_ssm_state: torch.Tensor):
|
|
self.ssm_states[layer_idx] = new_ssm_state.to(self.ssm_states[layer_idx].device)
|
|
return self.ssm_states[layer_idx]
|
|
|
|
def reset(self):
|
|
for layer_idx in range(len(self.conv_states)):
|
|
# In-place ops prevent breaking the static address
|
|
self.conv_states[layer_idx].zero_()
|
|
self.ssm_states[layer_idx].zero_()
|
|
|
|
class HQQQuantizedCache(QuantizedCache):
|
|
def __init__(
|
|
self,
|
|
config: PretrainedConfig,
|
|
nbits: int = 4,
|
|
axis_key: int = 0,
|
|
axis_value: int = 0,
|
|
q_group_size: int = 64,
|
|
residual_length: int = 128,
|
|
):
|
|
super().__init__("hqq", config, nbits, axis_key, axis_value, q_group_size, residual_length)
|
|
|
|
class HybridCache(StaticCache):
|
|
def __init__(self, config: PretrainedConfig, max_cache_len: int, *args, **kwargs):
|
|
super().__init__(config=config, max_cache_len=max_cache_len)
|
|
|
|
class HybridChunkedCache(StaticCache):
|
|
def __init__(self, config: PretrainedConfig, max_cache_len: int, *args, **kwargs):
|
|
super().__init__(config=config, max_cache_len=max_cache_len)
|
|
|
|
class OffloadedHybridCache(StaticCache):
|
|
def __init__(self, config: PretrainedConfig, max_cache_len: int, *args, **kwargs):
|
|
super().__init__(config=config, max_cache_len=max_cache_len, offloading=True)
|
|
|
|
class OffloadedStaticCache(StaticCache):
|
|
def __init__(self, config: PretrainedConfig, max_cache_len: int, *args, **kwargs):
|
|
super().__init__(config=config, max_cache_len=max_cache_len, offloading=True)
|
|
|
|
class OffloadedCache(DynamicCache):
|
|
def __init__(self) -> None:
|
|
super().__init__(offloading=True)
|
|
|
|
class QuantoQuantizedCache(QuantizedCache):
|
|
def __init__(
|
|
self,
|
|
config: PretrainedConfig,
|
|
nbits: int = 4,
|
|
axis_key: int = 0,
|
|
axis_value: int = 0,
|
|
q_group_size: int = 64,
|
|
residual_length: int = 128,
|
|
):
|
|
super().__init__("quanto", config, nbits, axis_key, axis_value, q_group_size, residual_length)
|
|
|
|
class SlidingWindowCache(StaticCache):
|
|
def __init__(self, config: PretrainedConfig, max_cache_len: int, *args, **kwargs):
|
|
super().__init__(config=config, max_cache_len=max_cache_len)
|
|
|
|
|
|
NEED_SETUP_CACHE_CLASSES_MAPPING = {}
|
|
QUANT_BACKEND_CLASSES_MAPPING = {}
|
|
CACHE_CONFIG_MAPPING = {}
|
|
|
|
if is_torch_available():
|
|
CACHE_CONFIG_MAPPING["quantized"] = QuantizedCacheConfig
|
|
CACHE_CONFIG_MAPPING["static"] = StaticCacheConfig
|
|
NEED_SETUP_CACHE_CLASSES_MAPPING = {
|
|
"static": StaticCache,
|
|
"offloaded_static": OffloadedStaticCache,
|
|
"sliding_window": SlidingWindowCache,
|
|
"hybrid": HybridCache,
|
|
"hybrid_chunked": HybridChunkedCache,
|
|
"offloaded_hybrid": OffloadedHybridCache,
|
|
"offloaded_hybrid_chunked": OffloadedHybridCache,
|
|
"mamba": MambaCache,
|
|
}
|
|
QUANT_BACKEND_CLASSES_MAPPING = {"quanto": QuantoQuantizedCache, "HQQ": HQQQuantizedCache}
|
|
ALL_CACHE_IMPLEMENTATIONS = (
|
|
list(NEED_SETUP_CACHE_CLASSES_MAPPING.keys()) + list(CACHE_CONFIG_MAPPING.keys()) + ["offloaded", "dynamic"]
|
|
) |