"`):
+ The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
+ token instead.
+ prefix_token (`str`, *optional*, defaults to `"▁"`):
+ Prefix token used for infilling.
+ suffix_token (`str`, *optional*, defaults to `"▁"`):
+ Suffix token used for infilling.
+ middle_token (`str`, *optional*, defaults to `"▁"`):
+ Middle token used for infilling.
+ eot_token (`str`, *optional*, defaults to `"▁"`):
+ End of text token used for infilling.
+ fill_token (`str`, *optional*, defaults to `""`):
+ The token used to split the input between the prefix and suffix.
+ suffix_first (`bool`, *optional*, default to `False`):
+ Whether the input prompt and suffix should be formatted with the suffix first.
+ additional_special_tokens (`List[str]`, *optional*):
+ Additional special tokens used by the tokenizer.
+ sp_model_kwargs (`dict`, *optional*):
+ Will be passed to the `SentencePieceProcessor.__init__()` method. The [Python wrapper for
+ SentencePiece](https://github.com/google/sentencepiece/tree/master/python) can be used, among other things,
+ to set:
+
+ - `enable_sampling`: Enable subword regularization.
+ - `nbest_size`: Sampling parameters for unigram. Invalid for BPE-Dropout.
+
+ - `nbest_size = {0,1}`: No sampling is performed.
+ - `nbest_size > 1`: samples from the nbest_size results.
+ - `nbest_size < 0`: assuming that nbest_size is infinite and samples from the all hypothesis (lattice)
+ using forward-filtering-and-backward-sampling algorithm.
+
+ - `alpha`: Smoothing parameter for unigram sampling, and dropout probability of merge operations for
+ BPE-dropout.
+ use_default_system_prompt (`bool`, *optional*, defaults to `False`):
+ Whether or not the default system prompt for Llama should be used.
"""
vocab_files_names = VOCAB_FILES_NAMES
@@ -106,99 +128,184 @@ def __init__(
unk_token="",
bos_token="",
eos_token="",
- pad_token=None,
+ prefix_token="▁",
+ middle_token="▁",
+ suffix_token="▁",
+ eot_token="▁",
+ fill_token="",
+ suffix_first=False,
sp_model_kwargs: Optional[Dict[str, Any]] = None,
add_bos_token=True,
add_eos_token=False,
clean_up_tokenization_spaces=False,
- legacy=True,
+ additional_special_tokens=None,
+ use_default_system_prompt=False,
**kwargs,
):
+ requires_backends(self, "protobuf")
self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs
bos_token = AddedToken(bos_token, lstrip=False, rstrip=False) if isinstance(bos_token, str) else bos_token
eos_token = AddedToken(eos_token, lstrip=False, rstrip=False) if isinstance(eos_token, str) else eos_token
unk_token = AddedToken(unk_token, lstrip=False, rstrip=False) if isinstance(unk_token, str) else unk_token
- pad_token = AddedToken(pad_token, lstrip=False, rstrip=False) if isinstance(pad_token, str) else pad_token
+
+ self.use_default_system_prompt = use_default_system_prompt
+ # mark tokens special to skip them
+ additional_special_tokens = additional_special_tokens or []
+ for token in [prefix_token, middle_token, suffix_token, eot_token]:
+ additional_special_tokens += [token] if token is not None else []
+
super().__init__(
bos_token=bos_token,
eos_token=eos_token,
unk_token=unk_token,
- pad_token=pad_token,
add_bos_token=add_bos_token,
add_eos_token=add_eos_token,
+ prefix_token=prefix_token,
+ middle_token=middle_token,
+ suffix_token=suffix_token,
+ eot_token=eot_token,
+ fill_token=fill_token,
sp_model_kwargs=self.sp_model_kwargs,
+ suffix_first=suffix_first,
clean_up_tokenization_spaces=clean_up_tokenization_spaces,
- legacy=legacy,
+ additional_special_tokens=additional_special_tokens,
+ use_default_system_prompt=use_default_system_prompt,
**kwargs,
)
- if legacy:
- logger.warning_once(
- f"You are using the legacy behaviour of the {self.__class__}. This means that tokens that come after special tokens will not be properly handled. We recommend you to"
- " read the related pull request available at https://github.com/huggingface/transformers/pull/24565"
- )
- self.legacy = legacy
self.vocab_file = vocab_file
self.add_bos_token = add_bos_token
self.add_eos_token = add_eos_token
- self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
- self.sp_model.Load(vocab_file)
+ self._prefix_token = prefix_token
+ self._middle_token = middle_token
+ self._suffix_token = suffix_token
+ self._eot_token = eot_token
+ self.fill_token = fill_token
+ self.suffix_first = suffix_first
+ self.sp_model = self.get_spm_processor()
- def __getstate__(self):
- state = self.__dict__.copy()
- state["sp_model"] = None
- state["sp_model_proto"] = self.sp_model.serialized_model_proto()
- return state
+ @property
+ def unk_token_length(self):
+ return len(self.sp_model.encode(str(self.unk_token)))
+
+ def get_spm_processor(self):
+ tokenizer = spm.SentencePieceProcessor(**self.sp_model_kwargs)
+ with open(self.vocab_file, "rb") as f:
+ sp_model = f.read()
+ model_pb2 = import_protobuf()
+ model = model_pb2.ModelProto.FromString(sp_model)
+ normalizer_spec = model_pb2.NormalizerSpec()
+ normalizer_spec.add_dummy_prefix = False
+ model.normalizer_spec.MergeFrom(normalizer_spec)
+ sp_model = model.SerializeToString()
+ tokenizer.LoadFromSerializedProto(sp_model)
+ return tokenizer
- def __setstate__(self, d):
- self.__dict__ = d
- self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
- self.sp_model.LoadFromSerializedProto(self.sp_model_proto)
+ @property
+ def prefix_token(self):
+ return self._prefix_token
+
+ @property
+ def prefix_id(self):
+ if self._prefix_token is None:
+ return None
+ return self.convert_tokens_to_ids(self.prefix_token)
+
+ @property
+ def middle_token(self):
+ return self._middle_token
+
+ @property
+ def middle_id(self):
+ if self._middle_token is None:
+ return None
+ return self.convert_tokens_to_ids(self.middle_token)
+
+ @property
+ def suffix_token(self):
+ return self._suffix_token
+
+ @property
+ def suffix_id(self):
+ if self._suffix_token is None:
+ return None
+ return self.convert_tokens_to_ids(self.suffix_token)
+
+ @property
+ def eot_token(self):
+ return self._eot_token
+
+ @property
+ def eot_id(self):
+ if self._eot_token is None:
+ return None
+ return self.convert_tokens_to_ids(self.eot_token)
@property
def vocab_size(self):
"""Returns vocab size"""
return self.sp_model.get_piece_size()
+ # Copied from transformers.models.llama.tokenization_llama.LlamaTokenizer.get_vocab
def get_vocab(self):
"""Returns vocab as a dict"""
vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)}
vocab.update(self.added_tokens_encoder)
return vocab
- # Copied from transformers.models.t5.tokenization_t5.T5Tokenizer.tokenize
- def tokenize(self, text, **kwargs) -> List[str]:
- # Replace the SPIECE_UNDERLINE with a space to make sure SPIECE_UNDERLINE is only used at
- # the beginning of the text
- if not self.legacy:
- text = SPIECE_UNDERLINE + text.replace(SPIECE_UNDERLINE, " ")
- return super().tokenize(text, **kwargs)
+ def tokenize(self, prefix, suffix=None, suffix_first=False, **kwargs) -> List[int]:
+ # add a prefix space to `prefix`
+ if self.fill_token is not None and self.fill_token in prefix and suffix is None:
+ prefix, suffix = prefix.split(self.fill_token)
+
+ if len(prefix) > 0:
+ prefix = SPIECE_UNDERLINE + prefix.replace(SPIECE_UNDERLINE, " ")
+
+ if suffix is None or len(suffix) < 1:
+ tokens = super().tokenize(prefix, **kwargs)
+ if len(tokens) > 1 and tokens[0] == SPIECE_UNDERLINE and tokens[1] in self.all_special_tokens:
+ tokens = tokens[1:]
+ return tokens
+
+ prefix_tokens = self._tokenize(prefix) # prefix has an extra `SPIECE_UNDERLINE`
+
+ if None in (self.prefix_id, self.middle_id, self.suffix_id):
+ raise ValueError(
+ "The input either includes a `prefix` and a `suffix` used for the infilling task,"
+ f" or can be split on the {self.fill_token} token, creating a suffix and prefix,"
+ " but the model does not support `infilling`."
+ )
+ suffix_tokens = self._tokenize(suffix) # make sure CodeLlama sp model does not mess up
- # Copied from transformers.models.t5.tokenization_t5.T5Tokenizer._tokenize
- def _tokenize(self, text):
+ suffix_first = suffix_first if suffix_first is not None else self.suffix_first
+ if suffix_first:
+ # format as " {suf} {pre}"
+ return [self.prefix_token, self.suffix_token] + suffix_tokens + [self.middle_token] + prefix_tokens
+ else:
+ # format as " {pre} {suf} "
+ return [self.prefix_token] + prefix_tokens + [self.suffix_token] + suffix_tokens + [self.middle_token]
+
+ def _tokenize(self, text, **kwargs):
"""
Returns a tokenized string.
- Since the sentencepiece internal model always adds a SPIECE_UNDERLINE, at the beginning of the provided text,
- we need to remove it by hand when the current text is a subsequence. This happens whenever the `self.tokenize`
- function is called with specials tokens: the input is split on the special tokens, and each subsequence is
- passed to `_tokenize`. Thus if a subsequence did not start with a `" "` or SPIECE_UNDERLINE, we have to remove
- the extra `SPIECE_UNDERLINE` prepended.
+ We de-activated the `add_dummy_prefix` option, thus the sentencepiece internals will always strip any
+ SPIECE_UNDERLINE. For example: `self.sp_model.encode(f"{SPIECE_UNDERLINE}Hey", out_type = str)` will give
+ `['H', 'e', 'y']` instead of `['▁He', 'y']`. Thus we always encode `f"{unk_token}text"` and strip the
+ `unk_token`. Here is an example with `unk_token = ""` and `unk_token_length = 4`.
+ `self.tokenizer.sp_model.encode(" Hey", out_type = str)[4:]`.
"""
- if not self.legacy:
- is_first = text.startswith(SPIECE_UNDERLINE)
- if is_first:
- text = text[1:]
-
tokens = self.sp_model.encode(text, out_type=str)
+ # 1. Encode string + prefix ex: " Hey"
+ tokens = self.sp_model.encode(self.unk_token + text, out_type=str)
+ # 2. Remove self.unk_token from ['<','unk','>', '▁Hey']
+ return tokens[self.unk_token_length :] if len(tokens) >= self.unk_token_length else tokens
- if not self.legacy and not is_first and not text.startswith(" ") and tokens[0].startswith(SPIECE_UNDERLINE):
- tokens = ([tokens[0][1:]] if len(tokens[0]) > 1 else []) + tokens[1:]
- return tokens
-
+ # Copied from transformers.models.llama.tokenization_llama.LlamaTokenizer._convert_token_to_id
def _convert_token_to_id(self, token):
"""Converts a token (str) in an id using the vocab."""
return self.sp_model.piece_to_id(token)
+ # Copied from transformers.models.llama.tokenization_llama.LlamaTokenizer._convert_id_to_token
def _convert_id_to_token(self, index):
"""Converts an index (integer) in a token (str) using the vocab."""
token = self.sp_model.IdToPiece(index)
@@ -206,23 +313,23 @@ def _convert_id_to_token(self, index):
def convert_tokens_to_string(self, tokens):
"""Converts a sequence of tokens (string) in a single string."""
+ # since we manually add the prefix space, we have to remove it when decoding
+ if tokens[0].startswith(SPIECE_UNDERLINE):
+ tokens[0] = tokens[0][1:]
+
current_sub_tokens = []
out_string = ""
- prev_is_special = False
- for i, token in enumerate(tokens):
+ for _, token in enumerate(tokens):
# make sure that special tokens are not decoded using sentencepiece model
if token in self.all_special_tokens:
- if not prev_is_special and i != 0:
- out_string += " "
out_string += self.sp_model.decode(current_sub_tokens) + token
- prev_is_special = True
current_sub_tokens = []
else:
current_sub_tokens.append(token)
- prev_is_special = False
out_string += self.sp_model.decode(current_sub_tokens)
return out_string
+ # Copied from transformers.models.llama.tokenization_llama.LlamaTokenizer.save_vocabulary
def save_vocabulary(self, save_directory, filename_prefix: Optional[str] = None) -> Tuple[str]:
"""
Save the vocabulary and special tokens file to a directory.
@@ -250,6 +357,7 @@ def save_vocabulary(self, save_directory, filename_prefix: Optional[str] = None)
return (out_vocab_file,)
+ # Copied from transformers.models.llama.tokenization_llama.LlamaTokenizer.build_inputs_with_special_tokens
def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
bos_token_id = [self.bos_token_id] if self.add_bos_token else []
eos_token_id = [self.eos_token_id] if self.add_eos_token else []
@@ -261,6 +369,7 @@ def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
return output
+ # Copied from transformers.models.llama.tokenization_llama.LlamaTokenizer.get_special_tokens_mask
def get_special_tokens_mask(
self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
) -> List[int]:
@@ -298,6 +407,7 @@ def get_special_tokens_mask(
+ eos_token_id
)
+ # Copied from transformers.models.llama.tokenization_llama.LlamaTokenizer.create_token_type_ids_from_sequences
def create_token_type_ids_from_sequences(
self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
) -> List[int]:
@@ -332,7 +442,7 @@ def create_token_type_ids_from_sequences(
return output
def _build_conversation_input_ids(self, conversation: "Conversation") -> List[int]:
- """Builds the input ids for a conversation.
+ r"""Builds the input ids for a conversation.
This is the format used in the provided examples. System prompts should be manually added at the beginning of
the conversation. If no system prompt is given, the `DEFAULT_SYSTEM_PROMPT` will be used.
```
@@ -346,8 +456,8 @@ def _build_conversation_input_ids(self, conversation: "Conversation") -> List[in
>>> from transformers import Conversation
>>> Conversation(
- ... "<>\n Only answer with emojis, and charades\n<>\n\nHow can I build a house in 10 septs?"
- ... )
+ ... "<>\n Complete the functions without any documentation\n<>\n\n `def remove_non_ascii(s: str) -> str:`"
+ ... ) # doctest: +IGNORE_RESULT
```
Args:
conversation (`Conversation`):
@@ -356,6 +466,21 @@ def _build_conversation_input_ids(self, conversation: "Conversation") -> List[in
`List[int]`:
Input ids for the conversation.
"""
+ if self.use_default_system_prompt:
+ if len(conversation.past_user_inputs) > 0:
+ if (
+ not conversation.past_user_inputs[0].startswith(B_SYS)
+ or E_SYS not in conversation.past_user_inputs[0]
+ ):
+ conversation.past_user_inputs[0] = (
+ B_SYS + DEFAULT_SYSTEM_PROMPT + E_SYS + conversation.past_user_inputs[0]
+ )
+ elif conversation.new_user_input:
+ if not conversation.new_user_input.startswith(B_SYS) or E_SYS not in conversation.new_user_input:
+ conversation.new_user_input = B_SYS + DEFAULT_SYSTEM_PROMPT + E_SYS + conversation.new_user_input
+ else:
+ raise ValueError("Last message must be from user")
+
dialogue = list(conversation.iter_texts())
if not all([is_user for is_user, msg in dialogue[::2]]) or not all(
[not is_user for is_user, msg in dialogue[1::2]]
@@ -365,14 +490,6 @@ def _build_conversation_input_ids(self, conversation: "Conversation") -> List[in
)
dialog_tokens: List[int] = []
- if len(conversation.past_user_inputs) > 0:
- if not conversation.past_user_inputs[0].startswith(B_SYS) or E_SYS not in conversation.past_user_inputs[0]:
- conversation.past_user_inputs[0] = (
- B_SYS + DEFAULT_SYSTEM_PROMPT + E_SYS + conversation.past_user_inputs[0]
- )
- elif not dialogue[0][1].startswith(B_SYS) or E_SYS not in dialogue[0][1]:
- dialogue[0] = (dialogue[0][0], B_SYS + DEFAULT_SYSTEM_PROMPT + E_SYS + dialogue[0][1])
-
dialog_tokens += sum(
[
[self.bos_token_id]
@@ -384,9 +501,18 @@ def _build_conversation_input_ids(self, conversation: "Conversation") -> List[in
],
[],
)
- if not (dialogue[-1][0]):
- raise ValueError(f"Last message must be from user, got {dialogue[-1]['role']}")
dialog_tokens += [self.bos_token_id] + self.encode(
f"{B_INST} {(dialogue[-1][1]).strip()} {E_INST}", add_special_tokens=False
)
return dialog_tokens
+
+ def __getstate__(self):
+ state = self.__dict__.copy()
+ state["sp_model"] = None
+ state["sp_model_proto"] = self.sp_model.serialized_model_proto()
+ return state
+
+ def __setstate__(self, d):
+ self.__dict__ = d
+ self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
+ self.sp_model.LoadFromSerializedProto(self.sp_model_proto)
diff --git a/mftcoder_accelerate/src/model/code_llama/tokenization_code_llama_fast.py b/mftcoder_accelerate/src/model/code_llama/tokenization_code_llama_fast.py
new file mode 100644
index 0000000..b492b65
--- /dev/null
+++ b/mftcoder_accelerate/src/model/code_llama/tokenization_code_llama_fast.py
@@ -0,0 +1,428 @@
+# coding=utf-8
+# Copyright 2023 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+from shutil import copyfile
+from typing import TYPE_CHECKING, List, Optional, Tuple
+
+from tokenizers import normalizers, processors
+
+from transformers.tokenization_utils_fast import PreTrainedTokenizerFast
+from transformers.utils import is_sentencepiece_available, logging
+from transformers.utils.versions import require_version
+
+
+if TYPE_CHECKING:
+ from transformers.pipelines.conversational import Conversation
+
+require_version("tokenizers>=0.13.3")
+
+if is_sentencepiece_available():
+ from .tokenization_code_llama import CodeLlamaTokenizer
+else:
+ CodeLlamaTokenizer = None
+
+logger = logging.get_logger(__name__)
+VOCAB_FILES_NAMES = {"vocab_file": "tokenizer.model", "tokenizer_file": "tokenizer.json"}
+
+SPIECE_UNDERLINE = "▁"
+
+
+B_INST, E_INST = "[INST]", "[/INST]"
+B_SYS, E_SYS = "<>\n", "\n<>\n\n"
+
+# fmt: off
+DEFAULT_SYSTEM_PROMPT = """You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. Your \
+answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure\
+ that your responses are socially unbiased and positive in nature.
+
+If a question does not make any sense, or is not factually coherent, explain why instead of answering something not \
+correct. If you don't know the answer to a question, please don't share false information."""
+# fmt: on
+
+
+class CodeLlamaTokenizerFast(PreTrainedTokenizerFast):
+ """
+ Construct a Llama tokenizer. Based on byte-level Byte-Pair-Encoding.
+
+ This uses notably ByteFallback and no normalization.
+
+ ```python
+ >>> from transformers import CodeLlamaTokenizerFast
+
+ >>> tokenizer = CodeLlamaTokenizerFast.from_pretrained("hf-internal-testing/llama-tokenizer")
+ >>> tokenizer.encode("Hello this is a test")
+ [1, 15043, 445, 338, 263, 1243]
+ ```
+
+ If you want to change the `bos_token` or the `eos_token`, make sure to specify them when initializing the model, or
+ call `tokenizer.update_post_processor()` to make sure that the post-processing is correctly done (otherwise the
+ values of the first token and final token of an encoded sequence will not be correct). For more details, checkout
+ [post-processors] (https://huggingface.co/docs/tokenizers/api/post-processors) documentation.
+
+
+ This tokenizer inherits from [`PreTrainedTokenizerFast`] which contains most of the main methods. Users should
+ refer to this superclass for more information regarding those methods. The default configuration match that of
+ [codellama/CodeLlama-7b-Instruct-hf](https://huggingface.co/codellama/CodeLlama-7b-Instruct-hf/blob/main/tokenizer_config.json)
+ which supports prompt infilling.
+
+ Args:
+ vocab_file (`str`):
+ [SentencePiece](https://github.com/google/sentencepiece) file (generally has a .model extension) that
+ contains the vocabulary necessary to instantiate a tokenizer.
+ tokenizer_file (`str`):
+ [tokenizers](https://github.com/huggingface/tokenizers) file (generally has a .json extension) that
+ contains everything needed to load the tokenizer.
+ clean_up_tokenization_spaces (`str`, *optional*, defaults to `False`):
+ Wether to cleanup spaces after decoding, cleanup consists in removing potential artifacts like extra
+ spaces.
+ bos_token (`str`, *optional*, defaults to `""`):
+ The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token.
+ eos_token (`str`, *optional*, defaults to `""`):
+ The end of sequence token.
+ unk_token (`str`, *optional*, defaults to `""`):
+ The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
+ token instead.
+ prefix_token (`str`, *optional*, defaults to `"▁"`):
+ Prefix token used for infilling.
+ suffix_token (`str`, *optional*, defaults to `"▁"`):
+ Suffix token used for infilling.
+ middle_token (`str`, *optional*, defaults to `"▁"`):
+ Middle token used for infilling.
+ eot_token (`str`, *optional*, defaults to `"▁"`):
+ End of text token used for infilling.
+ fill_token (`str`, *optional*, defaults to `""`):
+ The token used to split the input between the prefix and suffix.
+ suffix_first (`bool`, *optional*, default to `False`):
+ Whether the input prompt and suffix should be formatted with the suffix first.
+ additional_special_tokens (`List[str]`, *optional*):
+ Additional special tokens used by the tokenizer.
+ use_default_system_prompt (`bool`, *optional*, defaults to `True`):
+ Whether or not the default system prompt for Llama should be used.
+ """
+
+ vocab_files_names = VOCAB_FILES_NAMES
+ slow_tokenizer_class = CodeLlamaTokenizer
+ padding_side = "left"
+ model_input_names = ["input_ids", "attention_mask"]
+
+ def __init__(
+ self,
+ vocab_file=None,
+ tokenizer_file=None,
+ clean_up_tokenization_spaces=False,
+ unk_token="",
+ bos_token="",
+ eos_token="",
+ prefix_token="▁",
+ middle_token="▁",
+ suffix_token="▁",
+ eot_token="▁",
+ fill_token="",
+ additional_special_tokens=None,
+ add_bos_token=True,
+ add_eos_token=False,
+ use_default_system_prompt=False,
+ **kwargs,
+ ):
+ # mark tokens special to skip them
+ additional_special_tokens = additional_special_tokens or []
+ for token in [prefix_token, middle_token, suffix_token, eot_token]:
+ additional_special_tokens += [token] if token is not None else []
+ self.use_default_system_prompt = use_default_system_prompt
+
+ super().__init__(
+ vocab_file=vocab_file,
+ tokenizer_file=tokenizer_file,
+ clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+ additional_special_tokens=additional_special_tokens,
+ unk_token=unk_token,
+ bos_token=bos_token,
+ eos_token=eos_token,
+ prefix_token=prefix_token,
+ middle_token=middle_token,
+ suffix_token=suffix_token,
+ eot_token=eot_token,
+ fill_token=fill_token,
+ use_default_system_prompt=use_default_system_prompt,
+ **kwargs,
+ )
+ self._add_bos_token = add_bos_token
+ self._add_eos_token = add_eos_token
+ self.update_post_processor()
+
+ self.vocab_file = vocab_file
+
+ self._prefix_token = prefix_token
+ self._middle_token = middle_token
+ self._suffix_token = suffix_token
+ self._eot_token = eot_token
+ self.fill_token = fill_token
+
+ @property
+ def can_save_slow_tokenizer(self) -> bool:
+ return os.path.isfile(self.vocab_file) if self.vocab_file else False
+
+ # Copied from transformers.models.llama.tokenization_llama_fast.LlamaTokenizerFast.update_post_processor
+ def update_post_processor(self):
+ """
+ Updates the underlying post processor with the current `bos_token` and `eos_token`.
+ """
+ bos = self.bos_token
+ bos_token_id = self.bos_token_id
+
+ eos = self.eos_token
+ eos_token_id = self.eos_token_id
+
+ single = f"{(bos+':0 ') * self.add_bos_token}$A:0{(' '+eos+':0') * self.add_eos_token}"
+ pair = f"{single}{(' '+bos+':1') * self.add_bos_token} $B:1{(' '+eos+':1') * self.add_eos_token}"
+
+ special_tokens = []
+ if self.add_bos_token:
+ special_tokens.append((bos, bos_token_id))
+ if self.add_eos_token:
+ special_tokens.append((eos, eos_token_id))
+ self._tokenizer.post_processor = processors.TemplateProcessing(
+ single=single, pair=pair, special_tokens=special_tokens
+ )
+
+ @property
+ def prefix_token(self):
+ return self._prefix_token
+
+ @property
+ def prefix_id(self):
+ if self._prefix_token is None:
+ return None
+ return self.convert_tokens_to_ids(self.prefix_token)
+
+ @property
+ def middle_token(self):
+ return self._middle_token
+
+ @property
+ def middle_id(self):
+ if self._middle_token is None:
+ return None
+ return self.convert_tokens_to_ids(self.middle_token)
+
+ @property
+ def suffix_token(self):
+ return self._suffix_token
+
+ @property
+ def suffix_id(self):
+ if self._suffix_token is None:
+ return None
+ return self.convert_tokens_to_ids(self.suffix_token)
+
+ @property
+ def eot_id(self):
+ if self._eot_token is None:
+ return None
+ return self.convert_tokens_to_ids(self.eot_token)
+
+ @property
+ def eot_token(self):
+ return self._eot_token
+
+ @property
+ def add_eos_token(self):
+ return self._add_eos_token
+
+ @property
+ def add_bos_token(self):
+ return self._add_bos_token
+
+ @add_eos_token.setter
+ def add_eos_token(self, value):
+ self._add_eos_token = value
+ self.update_post_processor()
+
+ @add_bos_token.setter
+ def add_bos_token(self, value):
+ self._add_bos_token = value
+ self.update_post_processor()
+
+ def set_infilling_processor(self, reset, suffix_first=False, add_special_tokens=True):
+ if reset:
+ self._tokenizer.normalizer = normalizers.Sequence(
+ [
+ normalizers.Prepend(prepend="▁"),
+ normalizers.Replace(pattern=" ", content="▁"),
+ ]
+ )
+ self.update_post_processor()
+
+ self._tokenizer.normalizer = normalizers.Replace(pattern=" ", content="▁")
+ pair = [self.bos_token] if self.add_bos_token and add_special_tokens else []
+ special_tokens = [(self.bos_token, self.bos_token_id)] if self.add_bos_token and add_special_tokens else []
+ if suffix_first:
+ # format as " {suf} {pre}"
+ pair += [self.prefix_token, self.suffix_token, "$A", self.middle_token, "$B"]
+ special_tokens += [
+ (self.prefix_token, self.prefix_id),
+ (self.suffix_token, self.suffix_id),
+ (self.middle_token, self.middle_id),
+ ]
+ else:
+ # format as " {pre} {suf} "
+ pair += [self.prefix_token, "$A", self.suffix_token, "$B", self.middle_token]
+ special_tokens += [
+ (self.prefix_token, self.prefix_id),
+ (self.suffix_token, self.suffix_id),
+ (self.middle_token, self.middle_id),
+ ]
+
+ if self.add_eos_token and add_special_tokens:
+ pair += [self.eos_token]
+ special_tokens += [(self.eos_token, self.eos_token_id)]
+ self._tokenizer.post_processor = processors.TemplateProcessing(
+ single="$A", pair=pair, special_tokens=special_tokens
+ )
+
+ def encode_plus(self, text, text_pair=None, suffix_first=False, add_special_tokens=True, **kwargs):
+ # hack to make sure the input is pre-process but outside rust
+ text_pair = kwargs.pop("suffix", text_pair)
+ if self.fill_token in text and text_pair is None:
+ text, text_pair = text.split(self.fill_token)
+
+ if text_pair is None or len(text_pair) < 1:
+ return super().encode_plus(text, text_pair, add_special_tokens=add_special_tokens, **kwargs)
+
+ if None in (self.prefix_id, self.middle_id, self.suffix_id):
+ raise ValueError(
+ "Then input includes a `prefix` and a `suffix` used for the infilling task,"
+ " the `prefix_id, middle_id, suffix_id` must all be initialized. Current"
+ f" values : {self.prefix_id, self.middle_id, self.suffix_id}"
+ )
+
+ self.set_infilling_processor(False, suffix_first=suffix_first, add_special_tokens=add_special_tokens)
+ tokens = super().encode_plus(" " + text, text_pair=text_pair, add_special_tokens=True, **kwargs)
+ self.set_infilling_processor(True)
+ return tokens
+
+ # Copied from transformers.models.llama.tokenization_llama_fast.LlamaTokenizerFast.save_vocabulary
+ def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
+ if not self.can_save_slow_tokenizer:
+ raise ValueError(
+ "Your fast tokenizer does not have the necessary information to save the vocabulary for a slow "
+ "tokenizer."
+ )
+
+ if not os.path.isdir(save_directory):
+ logger.error(f"Vocabulary path ({save_directory}) should be a directory")
+ return
+ out_vocab_file = os.path.join(
+ save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
+ )
+
+ if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file):
+ copyfile(self.vocab_file, out_vocab_file)
+
+ return (out_vocab_file,)
+
+ def build_inputs_with_special_tokens(
+ self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
+ ) -> List[int]:
+ """
+ Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
+ adding special tokens. The special tokens depend on calling set_lang.
+
+ An NLLB sequence has the following format, where `X` represents the sequence:
+
+ - `input_ids` (for encoder) `X [eos, src_lang_code]`
+ - `decoder_input_ids`: (for decoder) `X [eos, tgt_lang_code]`
+
+ BOS is never used. Pairs of sequences are not the expected use case, but they will be handled without a
+ separator.
+
+ Args:
+ token_ids_0 (`List[int]`):
+ List of IDs to which the special tokens will be added.
+ token_ids_1 (`List[int]`, *optional*):
+ Optional second list of IDs for sequence pairs.
+
+ Returns:
+ `List[int]`: list of [input IDs](../glossary#input-ids) with the appropriate special tokens.
+ """
+ if token_ids_1 is None:
+ return self.bos_token_id + token_ids_0 + self.eos_token_id
+ return self.bos_token_id + token_ids_0 + token_ids_1 + self.eos_token_id
+
+ # Copied from transformers.models.code_llama.tokenization_code_llama.CodeLlamaTokenizer._build_conversation_input_ids
+ def _build_conversation_input_ids(self, conversation: "Conversation") -> List[int]:
+ r"""Builds the input ids for a conversation.
+ This is the format used in the provided examples. System prompts should be manually added at the beginning of
+ the conversation. If no system prompt is given, the `DEFAULT_SYSTEM_PROMPT` will be used.
+ ```
+ [INST] B_SYS SytemPrompt E_SYS Prompt [/INST] Answer
+ [INST] Prompt [/INST] Answer
+ [INST] Prompt [/INST]
+ ```
+
+ If you want to use your own system prompt, make sure to use both `B_SYS` and `E_SYS` use the following:
+ ```python
+ >>> from transformers import Conversation
+
+ >>> Conversation(
+ ... "<>\n Complete the functions without any documentation\n<>\n\n `def remove_non_ascii(s: str) -> str:`"
+ ... ) # doctest: +IGNORE_RESULT
+ ```
+ Args:
+ conversation (`Conversation`):
+ Conversation to build input ids for.
+ Returns:
+ `List[int]`:
+ Input ids for the conversation.
+ """
+ if self.use_default_system_prompt:
+ if len(conversation.past_user_inputs) > 0:
+ if (
+ not conversation.past_user_inputs[0].startswith(B_SYS)
+ or E_SYS not in conversation.past_user_inputs[0]
+ ):
+ conversation.past_user_inputs[0] = (
+ B_SYS + DEFAULT_SYSTEM_PROMPT + E_SYS + conversation.past_user_inputs[0]
+ )
+ elif conversation.new_user_input:
+ if not conversation.new_user_input.startswith(B_SYS) or E_SYS not in conversation.new_user_input:
+ conversation.new_user_input = B_SYS + DEFAULT_SYSTEM_PROMPT + E_SYS + conversation.new_user_input
+ else:
+ raise ValueError("Last message must be from user")
+
+ dialogue = list(conversation.iter_texts())
+ if not all([is_user for is_user, msg in dialogue[::2]]) or not all(
+ [not is_user for is_user, msg in dialogue[1::2]]
+ ):
+ raise ValueError(
+ "The model only supports 'user' and 'assistant' roles, starting with user and alternating (u/a/u/a/u...)"
+ )
+
+ dialog_tokens: List[int] = []
+ dialog_tokens += sum(
+ [
+ [self.bos_token_id]
+ + self.encode(
+ f"{B_INST} {(prompt[1]).strip()} {E_INST} {(answer[1]).strip()} ", add_special_tokens=False
+ )
+ + [self.eos_token_id]
+ for prompt, answer in zip(dialogue[::2], dialogue[1::2])
+ ],
+ [],
+ )
+ dialog_tokens += [self.bos_token_id] + self.encode(
+ f"{B_INST} {(dialogue[-1][1]).strip()} {E_INST}", add_special_tokens=False
+ )
+ return dialog_tokens
diff --git a/mft_peft_hf/src/model/code_llama/tokenization_llama.py b/mftcoder_accelerate/src/model/code_llama/tokenization_llama.py
similarity index 92%
rename from mft_peft_hf/src/model/code_llama/tokenization_llama.py
rename to mftcoder_accelerate/src/model/code_llama/tokenization_llama.py
index fb2cfce..98360d3 100644
--- a/mft_peft_hf/src/model/code_llama/tokenization_llama.py
+++ b/mftcoder_accelerate/src/model/code_llama/tokenization_llama.py
@@ -18,7 +18,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-"""Tokenization classes for LLaMA."""
+"""4.33.1 Tokenization classes for LLaMA."""
import os
from shutil import copyfile
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
@@ -141,7 +141,7 @@ def __init__(
logger.warning_once(
f"You are using the default legacy behaviour of the {self.__class__}. If you see this, DO NOT PANIC! This is"
" expected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you."
- " If you want to use the new behaviour, set `legacy=True`. This should only be set if you understand what it"
+ " If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it"
" means, and thouroughly read the reason why this was added as explained in"
" https://github.com/huggingface/transformers/pull/24565"
)
@@ -154,19 +154,25 @@ def __init__(
self.use_default_system_prompt = use_default_system_prompt
self.sp_model = self.get_spm_processor()
- self.unk_token_length = len(self.sp_model.encode(str(self.unk_token)))
+
+ @property
+ def unk_token_length(self):
+ return len(self.sp_model.encode(str(self.unk_token)))
# Copied from transformers.models.t5.tokenization_t5.T5Tokenizer.get_spm_processor
def get_spm_processor(self):
tokenizer = spm.SentencePieceProcessor(**self.sp_model_kwargs)
+ if self.legacy: # no dependency on protobuf
+ tokenizer.Load(self.vocab_file)
+ return tokenizer
+
with open(self.vocab_file, "rb") as f:
sp_model = f.read()
- model_pb2 = import_protobuf()
+ model_pb2 = import_protobuf(f"The new behaviour of {self.__class__.__name__} (with `self.legacy = False`)")
model = model_pb2.ModelProto.FromString(sp_model)
- if not self.legacy:
- normalizer_spec = model_pb2.NormalizerSpec()
- normalizer_spec.add_dummy_prefix = False
- model.normalizer_spec.MergeFrom(normalizer_spec)
+ normalizer_spec = model_pb2.NormalizerSpec()
+ normalizer_spec.add_dummy_prefix = False
+ model.normalizer_spec.MergeFrom(normalizer_spec)
sp_model = model.SerializeToString()
tokenizer.LoadFromSerializedProto(sp_model)
return tokenizer
@@ -194,18 +200,17 @@ def get_vocab(self):
return vocab
# Copied from transformers.models.t5.tokenization_t5.T5Tokenizer.tokenize
- def tokenize(self, text: "TextInput", **kwargs) -> List[str]:
+ def tokenize(self, text: "TextInput", add_special_tokens=False, **kwargs) -> List[str]:
"""
Converts a string to a list of tokens. If `self.legacy` is set to `False`, a prefix token is added unless the
first token is special.
"""
- if self.legacy:
+ if self.legacy or len(text) == 0:
return super().tokenize(text, **kwargs)
- if len(text) > 0:
- tokens = super().tokenize(SPIECE_UNDERLINE + text.replace(SPIECE_UNDERLINE, " "), **kwargs)
+ tokens = super().tokenize(SPIECE_UNDERLINE + text.replace(SPIECE_UNDERLINE, " "), **kwargs)
- if tokens[0] == SPIECE_UNDERLINE and tokens[1] in self.all_special_tokens:
+ if len(tokens) > 1 and tokens[0] == SPIECE_UNDERLINE and tokens[1] in self.all_special_tokens:
tokens = tokens[1:]
return tokens
@@ -220,13 +225,14 @@ def _tokenize(self, text, **kwargs):
`unk_token`. Here is an example with `unk_token = ""` and `unk_token_length = 4`.
`self.tokenizer.sp_model.encode(" Hey", out_type = str)[4:]`.
"""
- if self.legacy:
- return self.sp_model.encode(text, out_type=str)
-
- unk_token_length = len(self.sp_model.encode(str(self.unk_token)))
- text = self.unk_token + text
tokens = self.sp_model.encode(text, out_type=str)
- return tokens[unk_token_length:]
+ if self.legacy or not text.startswith((SPIECE_UNDERLINE, " ")):
+ return tokens
+
+ # 1. Encode string + prefix ex: " Hey"
+ tokens = self.sp_model.encode(self.unk_token + text, out_type=str)
+ # 2. Remove self.unk_token from ['<','unk','>', '▁Hey']
+ return tokens[self.unk_token_length :] if len(tokens) >= self.unk_token_length else tokens
def _convert_token_to_id(self, token):
"""Converts a token (str) in an id using the vocab."""
diff --git a/mft_peft_hf/src/model/code_llama/tokenization_llama_fast.py b/mftcoder_accelerate/src/model/code_llama/tokenization_llama_fast.py
similarity index 98%
rename from mft_peft_hf/src/model/code_llama/tokenization_llama_fast.py
rename to mftcoder_accelerate/src/model/code_llama/tokenization_llama_fast.py
index f7ad2ec..53a4b0b 100644
--- a/mft_peft_hf/src/model/code_llama/tokenization_llama_fast.py
+++ b/mftcoder_accelerate/src/model/code_llama/tokenization_llama_fast.py
@@ -128,7 +128,10 @@ def __init__(
self.update_post_processor()
self.use_default_system_prompt = use_default_system_prompt
self.vocab_file = vocab_file
- self.can_save_slow_tokenizer = False if not self.vocab_file else True
+
+ @property
+ def can_save_slow_tokenizer(self) -> bool:
+ return os.path.isfile(self.vocab_file) if self.vocab_file else False
def update_post_processor(self):
"""
diff --git a/mftcoder_accelerate/src/model/deepseek_v2/configuration_deepseek.py b/mftcoder_accelerate/src/model/deepseek_v2/configuration_deepseek.py
new file mode 100644
index 0000000..82e0f5d
--- /dev/null
+++ b/mftcoder_accelerate/src/model/deepseek_v2/configuration_deepseek.py
@@ -0,0 +1,206 @@
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+DEEPSEEK_PRETRAINED_CONFIG_ARCHIVE_MAP = {}
+class DeepseekV2Config(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`DeepseekV2Model`]. It is used to instantiate an DeepSeek
+ model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
+ defaults will yield a similar configuration to that of the DeepSeek-V2.
+
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+ documentation from [`PretrainedConfig`] for more information.
+
+
+ Args:
+ vocab_size (`int`, *optional*, defaults to 102400):
+ Vocabulary size of the Deep model. Defines the number of different tokens that can be represented by the
+ `inputs_ids` passed when calling [`DeepseekV2Model`]
+ hidden_size (`int`, *optional*, defaults to 4096):
+ Dimension of the hidden representations.
+ intermediate_size (`int`, *optional*, defaults to 11008):
+ Dimension of the MLP representations.
+ moe_intermediate_size (`int`, *optional*, defaults to 1407):
+ Dimension of the MoE representations.
+ num_hidden_layers (`int`, *optional*, defaults to 32):
+ Number of hidden layers in the Transformer decoder.
+ num_attention_heads (`int`, *optional*, defaults to 32):
+ Number of attention heads for each attention layer in the Transformer decoder.
+ n_shared_experts (`int`, *optional*, defaults to None):
+ Number of shared experts, None means dense model.
+ n_routed_experts (`int`, *optional*, defaults to None):
+ Number of routed experts, None means dense model.
+ routed_scaling_factor (`float`, *optional*, defaults to 1.0):
+ Scaling factor or routed experts.
+ topk_method (`str`, *optional*, defaults to `gready`):
+ Topk method used in routed gate.
+ n_group (`int`, *optional*, defaults to None):
+ Number of groups for routed experts.
+ topk_group (`int`, *optional*, defaults to None):
+ Number of selected groups for each token(for each token, ensuring the selected experts is only within `topk_group` groups).
+ num_experts_per_tok (`int`, *optional*, defaults to None):
+ Number of selected experts, None means dense model.
+ moe_layer_freq (`int`, *optional*, defaults to 1):
+ The frequency of the MoE layer: one expert layer for every `moe_layer_freq - 1` dense layers.
+ first_k_dense_replace (`int`, *optional*, defaults to 0):
+ Number of dense layers in shallow layers(embed->dense->dense->...->dense->moe->moe...->lm_head).
+ \--k dense layers--/
+ norm_topk_prob (`bool`, *optional*, defaults to False):
+ Whether to normalize the weights of the routed experts.
+ scoring_func (`str`, *optional*, defaults to 'softmax'):
+ Method of computing expert weights.
+ aux_loss_alpha (`float`, *optional*, defaults to 0.001):
+ Auxiliary loss weight coefficient.
+ seq_aux = (`bool`, *optional*, defaults to True):
+ Whether to compute the auxiliary loss for each individual sample.
+ num_key_value_heads (`int`, *optional*):
+ This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+ `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
+ `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When
+ converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
+ by meanpooling all the original heads within that group. For more details checkout [this
+ paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to
+ `num_attention_heads`.
+ hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+ The non-linear activation function (function or string) in the decoder.
+ max_position_embeddings (`int`, *optional*, defaults to 2048):
+ The maximum sequence length that this model might ever be used with.
+ initializer_range (`float`, *optional*, defaults to 0.02):
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+ rms_norm_eps (`float`, *optional*, defaults to 1e-06):
+ The epsilon used by the rms normalization layers.
+ use_cache (`bool`, *optional*, defaults to `True`):
+ Whether or not the model should return the last key/values attentions (not used by all models). Only
+ relevant if `config.is_decoder=True`.
+ pad_token_id (`int`, *optional*):
+ Padding token id.
+ bos_token_id (`int`, *optional*, defaults to 1):
+ Beginning of stream token id.
+ eos_token_id (`int`, *optional*, defaults to 2):
+ End of stream token id.
+ pretraining_tp (`int`, *optional*, defaults to 1):
+ Experimental feature. Tensor parallelism rank used during pretraining. Please refer to [this
+ document](https://huggingface.co/docs/transformers/parallelism) to understand more about it. This value is
+ necessary to ensure exact reproducibility of the pretraining results. Please refer to [this
+ issue](https://github.com/pytorch/pytorch/issues/76232).
+ tie_word_embeddings (`bool`, *optional*, defaults to `False`):
+ Whether to tie weight embeddings
+ rope_theta (`float`, *optional*, defaults to 10000.0):
+ The base period of the RoPE embeddings.
+ rope_scaling (`Dict`, *optional*):
+ Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling
+ strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is
+ `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update
+ `max_position_embeddings` to the expected new maximum.
+ attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`):
+ Whether to use a bias in the query, key, value and output projection layers during self-attention.
+ attention_dropout (`float`, *optional*, defaults to 0.0):
+ The dropout ratio for the attention probabilities.
+
+ ```python
+ >>> from transformers import DeepseekV2Model, DeepseekV2Config
+
+ >>> # Initializing a Deepseek-V2 style configuration
+ >>> configuration = DeepseekV2Config()
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "deepseek_v2"
+ keys_to_ignore_at_inference = ["past_key_values"]
+
+ def __init__(
+ self,
+ vocab_size=102400,
+ hidden_size=4096,
+ intermediate_size=11008,
+ moe_intermediate_size = 1407,
+ num_hidden_layers=30,
+ num_attention_heads=32,
+ num_key_value_heads=32,
+ n_shared_experts = None,
+ n_routed_experts = None,
+ ep_size = 1,
+ routed_scaling_factor = 1.0,
+ kv_lora_rank = 512,
+ q_lora_rank = 1536,
+ qk_rope_head_dim = 64,
+ v_head_dim = 128,
+ qk_nope_head_dim = 128,
+ topk_method = 'gready',
+ n_group = None,
+ topk_group = None,
+ num_experts_per_tok = None,
+ moe_layer_freq = 1,
+ first_k_dense_replace = 0,
+ norm_topk_prob = False,
+ scoring_func = 'softmax',
+ aux_loss_alpha = 0.001,
+ seq_aux = True,
+ hidden_act="silu",
+ max_position_embeddings=2048,
+ initializer_range=0.02,
+ rms_norm_eps=1e-6,
+ use_cache=True,
+ pad_token_id=None,
+ bos_token_id=100000,
+ eos_token_id=100001,
+ pretraining_tp=1,
+ tie_word_embeddings=False,
+ rope_theta=10000.0,
+ rope_scaling=None,
+ attention_bias=False,
+ attention_dropout=0.0,
+ **kwargs,
+ ):
+ self.vocab_size = vocab_size
+ self.max_position_embeddings = max_position_embeddings
+ self.hidden_size = hidden_size
+ self.intermediate_size = intermediate_size
+ self.moe_intermediate_size = moe_intermediate_size
+ self.num_hidden_layers = num_hidden_layers
+ self.num_attention_heads = num_attention_heads
+ self.n_shared_experts = n_shared_experts
+ self.n_routed_experts = n_routed_experts
+ self.ep_size = ep_size
+ self.routed_scaling_factor = routed_scaling_factor
+ self.kv_lora_rank = kv_lora_rank
+ self.q_lora_rank = q_lora_rank
+ self.qk_rope_head_dim = qk_rope_head_dim
+ self.v_head_dim = v_head_dim
+ self.qk_nope_head_dim = qk_nope_head_dim
+ self.topk_method = topk_method
+ self.n_group = n_group
+ self.topk_group = topk_group
+ self.num_experts_per_tok = num_experts_per_tok
+ self.moe_layer_freq = moe_layer_freq
+ self.first_k_dense_replace = first_k_dense_replace
+ self.norm_topk_prob = norm_topk_prob
+ self.scoring_func = scoring_func
+ self.aux_loss_alpha = aux_loss_alpha
+ self.seq_aux = seq_aux
+ # for backward compatibility
+ if num_key_value_heads is None:
+ num_key_value_heads = num_attention_heads
+
+ self.num_key_value_heads = num_key_value_heads
+ self.hidden_act = hidden_act
+ self.initializer_range = initializer_range
+ self.rms_norm_eps = rms_norm_eps
+ self.pretraining_tp = pretraining_tp
+ self.use_cache = use_cache
+ self.rope_theta = rope_theta
+ self.rope_scaling = rope_scaling
+ self.attention_bias = attention_bias
+ self.attention_dropout = attention_dropout
+
+ super().__init__(
+ pad_token_id=pad_token_id,
+ bos_token_id=bos_token_id,
+ eos_token_id=eos_token_id,
+ tie_word_embeddings=tie_word_embeddings,
+ **kwargs,
+ )
\ No newline at end of file
diff --git a/mftcoder_accelerate/src/model/deepseek_v2/modeling_deepseek.py b/mftcoder_accelerate/src/model/deepseek_v2/modeling_deepseek.py
new file mode 100644
index 0000000..d1d5e88
--- /dev/null
+++ b/mftcoder_accelerate/src/model/deepseek_v2/modeling_deepseek.py
@@ -0,0 +1,1925 @@
+# coding=utf-8
+# Copyright 2023 DeepSeek-AI and The HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" PyTorch DeepSeek model."""
+import math
+import warnings
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.modeling_attn_mask_utils import (
+ AttentionMaskConverter,
+ _prepare_4d_attention_mask,
+ _prepare_4d_causal_attention_mask,
+)
+from transformers.modeling_outputs import (
+ BaseModelOutputWithPast,
+ CausalLMOutputWithPast,
+ SequenceClassifierOutputWithPast,
+)
+from transformers.modeling_utils import PreTrainedModel
+from transformers.pytorch_utils import (
+ ALL_LAYERNORM_LAYERS,
+ is_torch_greater_or_equal_than_1_13,
+)
+from transformers.utils import (
+ add_start_docstrings,
+ add_start_docstrings_to_model_forward,
+ is_flash_attn_2_available,
+ is_flash_attn_greater_or_equal_2_10,
+ logging,
+ replace_return_docstrings,
+)
+from transformers.utils.import_utils import is_torch_fx_available
+from .configuration_deepseek import DeepseekV2Config
+import torch.distributed as dist
+import numpy as np
+
+if is_flash_attn_2_available():
+ from flash_attn import flash_attn_func, flash_attn_varlen_func
+ from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa
+
+
+# This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph.
+# It means that the function will not be traced through and simply appear as a node in the graph.
+if is_torch_fx_available():
+ if not is_torch_greater_or_equal_than_1_13:
+ import torch.fx
+
+ _prepare_4d_causal_attention_mask = torch.fx.wrap(_prepare_4d_causal_attention_mask)
+
+
+logger = logging.get_logger(__name__)
+
+_CONFIG_FOR_DOC = "DeepseekV2Config"
+
+
+def _get_unpad_data(attention_mask):
+ seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
+ indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
+ max_seqlen_in_batch = seqlens_in_batch.max().item()
+ cu_seqlens = F.pad(
+ torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0)
+ )
+ return (
+ indices,
+ cu_seqlens,
+ max_seqlen_in_batch,
+ )
+
+
+class DeepseekV2RMSNorm(nn.Module):
+ def __init__(self, hidden_size, eps=1e-6):
+ """
+ DeepseekV2RMSNorm is equivalent to T5LayerNorm
+ """
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(hidden_size))
+ self.variance_epsilon = eps
+
+ def forward(self, hidden_states):
+ input_dtype = hidden_states.dtype
+ hidden_states = hidden_states.to(torch.float32)
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+ return self.weight * hidden_states.to(input_dtype)
+
+
+ALL_LAYERNORM_LAYERS.append(DeepseekV2RMSNorm)
+
+
+class DeepseekV2RotaryEmbedding(nn.Module):
+ def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
+ super().__init__()
+
+ self.dim = dim
+ self.max_position_embeddings = max_position_embeddings
+ self.base = base
+ inv_freq = 1.0 / (
+ self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim)
+ )
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+ # Build here to make `torch.jit.trace` work.
+ self._set_cos_sin_cache(
+ seq_len=max_position_embeddings,
+ device=self.inv_freq.device,
+ dtype=torch.get_default_dtype(),
+ )
+ self.max_seq_len_cached = None
+
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
+ self.max_seq_len_cached = seq_len
+ t = torch.arange(
+ self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype
+ )
+
+ freqs = torch.outer(t, self.inv_freq.to(t.device))
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
+ emb = torch.cat((freqs, freqs), dim=-1)
+ self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
+ self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
+
+ def forward(self, x, seq_len=None):
+ # x: [bs, num_attention_heads, seq_len, head_size]
+ if self.max_seq_len_cached is None or seq_len > self.max_seq_len_cached:
+ self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
+
+ return (
+ self.cos_cached[:seq_len].to(dtype=x.dtype),
+ self.sin_cached[:seq_len].to(dtype=x.dtype),
+ )
+
+
+# Copied from transformers.models.llama.modeling_llama.LlamaLinearScalingRotaryEmbedding with Llama->DeepseekV2
+class DeepseekV2LinearScalingRotaryEmbedding(DeepseekV2RotaryEmbedding):
+ """DeepseekV2RotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""
+
+ def __init__(
+ self,
+ dim,
+ max_position_embeddings=2048,
+ base=10000,
+ device=None,
+ scaling_factor=1.0,
+ ):
+ self.scaling_factor = scaling_factor
+ super().__init__(dim, max_position_embeddings, base, device)
+
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
+ self.max_seq_len_cached = seq_len
+ t = torch.arange(
+ self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype
+ )
+ t = t / self.scaling_factor
+
+ freqs = torch.outer(t, self.inv_freq)
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
+ emb = torch.cat((freqs, freqs), dim=-1)
+ self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
+ self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
+
+
+# Copied from transformers.models.llama.modeling_llama.LlamaDynamicNTKScalingRotaryEmbedding with Llama->DeepseekV2
+class DeepseekV2DynamicNTKScalingRotaryEmbedding(DeepseekV2RotaryEmbedding):
+ """DeepseekV2RotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""
+
+ def __init__(
+ self,
+ dim,
+ max_position_embeddings=2048,
+ base=10000,
+ device=None,
+ scaling_factor=1.0,
+ ):
+ self.scaling_factor = scaling_factor
+ super().__init__(dim, max_position_embeddings, base, device)
+
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
+ self.max_seq_len_cached = seq_len
+
+ if seq_len > self.max_position_embeddings:
+ base = self.base * (
+ (self.scaling_factor * seq_len / self.max_position_embeddings)
+ - (self.scaling_factor - 1)
+ ) ** (self.dim / (self.dim - 2))
+ inv_freq = 1.0 / (
+ base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim)
+ )
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+ t = torch.arange(
+ self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype
+ )
+
+ freqs = torch.outer(t, self.inv_freq)
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
+ emb = torch.cat((freqs, freqs), dim=-1)
+ self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
+ self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
+
+
+# Inverse dim formula to find dim based on number of rotations
+def yarn_find_correction_dim(
+ num_rotations, dim, base=10000, max_position_embeddings=2048
+):
+ return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / (
+ 2 * math.log(base)
+ )
+
+
+# Find dim range bounds based on rotations
+def yarn_find_correction_range(
+ low_rot, high_rot, dim, base=10000, max_position_embeddings=2048
+):
+ low = math.floor(
+ yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings)
+ )
+ high = math.ceil(
+ yarn_find_correction_dim(high_rot, dim, base, max_position_embeddings)
+ )
+ return max(low, 0), min(high, dim - 1) # Clamp values just in case
+
+
+def yarn_get_mscale(scale=1, mscale=1):
+ if scale <= 1:
+ return 1.0
+ return 0.1 * mscale * math.log(scale) + 1.0
+
+
+def yarn_linear_ramp_mask(min, max, dim):
+ if min == max:
+ max += 0.001 # Prevent singularity
+
+ linear_func = (torch.arange(dim, dtype=torch.float32) - min) / (max - min)
+ ramp_func = torch.clamp(linear_func, 0, 1)
+ return ramp_func
+
+
+class DeepseekV2YarnRotaryEmbedding(DeepseekV2RotaryEmbedding):
+
+ def __init__(
+ self,
+ dim,
+ max_position_embeddings=2048,
+ base=10000,
+ device=None,
+ scaling_factor=1.0,
+ original_max_position_embeddings=4096,
+ beta_fast=32,
+ beta_slow=1,
+ mscale=1,
+ mscale_all_dim=0,
+ ):
+ self.scaling_factor = scaling_factor
+ self.original_max_position_embeddings = original_max_position_embeddings
+ self.beta_fast = beta_fast
+ self.beta_slow = beta_slow
+ self.mscale = mscale
+ self.mscale_all_dim = mscale_all_dim
+ super().__init__(dim, max_position_embeddings, base, device)
+
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
+ self.max_seq_len_cached = seq_len
+ dim = self.dim
+
+ freq_extra = 1.0 / (
+ self.base
+ ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim)
+ )
+ freq_inter = 1.0 / (
+ self.scaling_factor
+ * self.base
+ ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim)
+ )
+
+ low, high = yarn_find_correction_range(
+ self.beta_fast,
+ self.beta_slow,
+ dim,
+ self.base,
+ self.original_max_position_embeddings,
+ )
+ inv_freq_mask = 1.0 - yarn_linear_ramp_mask(low, high, dim // 2).to(
+ device=device, dtype=torch.float32
+ )
+ inv_freq = freq_inter * (1 - inv_freq_mask) + freq_extra * inv_freq_mask
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+ t = torch.arange(seq_len, device=device, dtype=torch.float32)
+
+ freqs = torch.outer(t, inv_freq)
+
+ _mscale = float(
+ yarn_get_mscale(self.scaling_factor, self.mscale)
+ / yarn_get_mscale(self.scaling_factor, self.mscale_all_dim)
+ )
+
+ emb = torch.cat((freqs, freqs), dim=-1)
+ self.register_buffer(
+ "cos_cached", (emb.cos() * _mscale).to(dtype), persistent=False
+ )
+ self.register_buffer(
+ "sin_cached", (emb.sin() * _mscale).to(dtype), persistent=False
+ )
+
+
+# Copied from transformers.models.llama.modeling_llama.rotate_half
+def rotate_half(x):
+ """Rotates half the hidden dims of the input."""
+ x1 = x[..., : x.shape[-1] // 2]
+ x2 = x[..., x.shape[-1] // 2 :]
+ return torch.cat((-x2, x1), dim=-1)
+
+
+# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
+ """Applies Rotary Position Embedding to the query and key tensors.
+
+ Args:
+ q (`torch.Tensor`): The query tensor.
+ k (`torch.Tensor`): The key tensor.
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
+ position_ids (`torch.Tensor`):
+ The position indices of the tokens corresponding to the query and key tensors. For example, this can be
+ used to pass offsetted position ids when working with a KV-cache.
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+ Returns:
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+ """
+ cos = cos[position_ids].unsqueeze(unsqueeze_dim)
+ sin = sin[position_ids].unsqueeze(unsqueeze_dim)
+
+ b, h, s, d = q.shape
+ q = q.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d)
+
+ b, h, s, d = k.shape
+ k = k.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d)
+
+ q_embed = (q * cos) + (rotate_half(q) * sin)
+ k_embed = (k * cos) + (rotate_half(k) * sin)
+ return q_embed, k_embed
+
+
+class DeepseekV2MLP(nn.Module):
+ def __init__(self, config, hidden_size=None, intermediate_size=None):
+ super().__init__()
+ self.config = config
+ self.hidden_size = config.hidden_size if hidden_size is None else hidden_size
+ self.intermediate_size = (
+ config.intermediate_size if intermediate_size is None else intermediate_size
+ )
+
+ self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+ self.act_fn = ACT2FN[config.hidden_act]
+
+ def forward(self, x):
+ down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+ return down_proj
+
+
+class MoEGate(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.top_k = config.num_experts_per_tok
+ self.n_routed_experts = config.n_routed_experts
+ self.routed_scaling_factor = config.routed_scaling_factor
+ self.scoring_func = config.scoring_func
+ self.alpha = config.aux_loss_alpha
+ self.seq_aux = config.seq_aux
+ self.topk_method = config.topk_method
+ self.n_group = config.n_group
+ self.topk_group = config.topk_group
+
+ # topk selection algorithm
+ self.norm_topk_prob = config.norm_topk_prob
+ self.gating_dim = config.hidden_size
+ self.weight = nn.Parameter(
+ torch.empty((self.n_routed_experts, self.gating_dim))
+ )
+ self.reset_parameters()
+
+ def reset_parameters(self) -> None:
+ import torch.nn.init as init
+
+ init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+
+ def forward(self, hidden_states):
+ bsz, seq_len, h = hidden_states.shape
+ ### compute gating score
+ hidden_states = hidden_states.view(-1, h)
+ logits = F.linear(
+ hidden_states.type(torch.float32), self.weight.type(torch.float32), None
+ )
+ if self.scoring_func == "softmax":
+ scores = logits.softmax(dim=-1, dtype=torch.float32)
+ else:
+ raise NotImplementedError(
+ f"insupportable scoring function for MoE gating: {self.scoring_func}"
+ )
+
+ ### select top-k experts
+ if self.topk_method == "greedy":
+ topk_weight, topk_idx = torch.topk(
+ scores, k=self.top_k, dim=-1, sorted=False
+ )
+ elif self.topk_method == "group_limited_greedy":
+ group_scores = (
+ scores.view(bsz * seq_len, self.n_group, -1).max(dim=-1).values
+ ) # [n, n_group]
+ group_idx = torch.topk(
+ group_scores, k=self.topk_group, dim=-1, sorted=False
+ )[
+ 1
+ ] # [n, top_k_group]
+ group_mask = torch.zeros_like(group_scores) # [n, n_group]
+ group_mask.scatter_(1, group_idx, 1) # [n, n_group]
+ score_mask = (
+ group_mask.unsqueeze(-1)
+ .expand(
+ bsz * seq_len, self.n_group, self.n_routed_experts // self.n_group
+ )
+ .reshape(bsz * seq_len, -1)
+ ) # [n, e]
+ tmp_scores = scores.masked_fill(~score_mask.bool(), 0.0) # [n, e]
+ topk_weight, topk_idx = torch.topk(
+ tmp_scores, k=self.top_k, dim=-1, sorted=False
+ )
+
+ ### norm gate to sum 1
+ if self.top_k > 1 and self.norm_topk_prob:
+ denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20
+ topk_weight = topk_weight / denominator
+ else:
+ topk_weight = topk_weight * self.routed_scaling_factor
+ ### expert-level computation auxiliary loss
+ if self.training and self.alpha > 0.0:
+ scores_for_aux = scores
+ aux_topk = self.top_k
+ # always compute aux loss based on the naive greedy topk method
+ topk_idx_for_aux_loss = topk_idx.view(bsz, -1)
+ if self.seq_aux:
+ scores_for_seq_aux = scores_for_aux.view(bsz, seq_len, -1)
+ ce = torch.zeros(
+ bsz, self.n_routed_experts, device=hidden_states.device
+ )
+ ce.scatter_add_(
+ 1,
+ topk_idx_for_aux_loss,
+ torch.ones(bsz, seq_len * aux_topk, device=hidden_states.device),
+ ).div_(seq_len * aux_topk / self.n_routed_experts)
+ aux_loss = (ce * scores_for_seq_aux.mean(dim=1)).sum(
+ dim=1
+ ).mean() * self.alpha
+ else:
+ mask_ce = F.one_hot(
+ topk_idx_for_aux_loss.view(-1), num_classes=self.n_routed_experts
+ )
+ ce = mask_ce.float().mean(0)
+ Pi = scores_for_aux.mean(0)
+ fi = ce * self.n_routed_experts
+ aux_loss = (Pi * fi).sum() * self.alpha
+ else:
+ aux_loss = None
+ return topk_idx, topk_weight, aux_loss
+
+
+class AddAuxiliaryLoss(torch.autograd.Function):
+ """
+ The trick function of adding auxiliary (aux) loss,
+ which includes the gradient of the aux loss during backpropagation.
+ """
+
+ @staticmethod
+ def forward(ctx, x, loss):
+ assert loss.numel() == 1
+ ctx.dtype = loss.dtype
+ ctx.required_aux_loss = loss.requires_grad
+ return x
+
+ @staticmethod
+ def backward(ctx, grad_output):
+ grad_loss = None
+ if ctx.required_aux_loss:
+ grad_loss = torch.ones(1, dtype=ctx.dtype, device=grad_output.device)
+ return grad_output, grad_loss
+
+
+class DeepseekV2MoE(nn.Module):
+ """
+ A mixed expert module containing shared experts.
+ """
+
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.num_experts_per_tok = config.num_experts_per_tok
+
+ if hasattr(config, "ep_size") and config.ep_size > 1:
+ assert config.ep_size == dist.get_world_size()
+ self.ep_size = config.ep_size
+ self.experts_per_rank = config.n_routed_experts // config.ep_size
+ self.ep_rank = dist.get_rank()
+ self.experts = nn.ModuleList(
+ [
+ (
+ DeepseekV2MLP(
+ config, intermediate_size=config.moe_intermediate_size
+ )
+ if i >= self.ep_rank * self.experts_per_rank
+ and i < (self.ep_rank + 1) * self.experts_per_rank
+ else None
+ )
+ for i in range(config.n_routed_experts)
+ ]
+ )
+ else:
+ self.ep_size = 1
+ self.experts_per_rank = config.n_routed_experts
+ self.ep_rank = 0
+ self.experts = nn.ModuleList(
+ [
+ DeepseekV2MLP(config, intermediate_size=config.moe_intermediate_size)
+ for i in range(config.n_routed_experts)
+ ]
+ )
+ self.gate = MoEGate(config)
+ if config.n_shared_experts is not None:
+ intermediate_size = config.moe_intermediate_size * config.n_shared_experts
+ self.shared_experts = DeepseekV2MLP(
+ config=config, intermediate_size=intermediate_size
+ )
+
+ def forward(self, hidden_states):
+ # save dtype before computation
+ input_dtype = hidden_states.dtype
+ identity = hidden_states
+ orig_shape = hidden_states.shape
+ topk_idx, topk_weight, aux_loss = self.gate(hidden_states)
+ hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
+ flat_topk_idx = topk_idx.view(-1)
+ if self.training:
+ hidden_states = hidden_states.repeat_interleave(
+ self.num_experts_per_tok, dim=0
+ )
+ y = torch.empty_like(hidden_states)
+ for i, expert in enumerate(self.experts):
+ y[flat_topk_idx == i] = expert(hidden_states[flat_topk_idx == i])
+ y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1)
+ y = y.view(*orig_shape)
+ y = AddAuxiliaryLoss.apply(y, aux_loss)
+ else:
+ y = self.moe_infer(hidden_states, topk_idx, topk_weight).view(*orig_shape)
+ if self.config.n_shared_experts is not None:
+ y = y + self.shared_experts(identity)
+ # keep dtype same after moe forward
+ return y.to(input_dtype)
+
+ @torch.no_grad()
+ def moe_infer(self, x, topk_ids, topk_weight):
+ cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts)))
+ cnts.scatter_(1, topk_ids, 1)
+ tokens_per_expert = cnts.sum(dim=0)
+ idxs = topk_ids.view(-1).argsort()
+ sorted_tokens = x[idxs // topk_ids.shape[1]]
+ sorted_tokens_shape = sorted_tokens.shape
+ if self.ep_size > 1:
+ tokens_per_ep_rank = tokens_per_expert.view(self.ep_size, -1).sum(dim=1)
+ tokens_per_expert_group = tokens_per_expert.new_empty(
+ tokens_per_expert.shape[0]
+ )
+ dist.all_to_all_single(tokens_per_expert_group, tokens_per_expert)
+ output_splits = (
+ tokens_per_expert_group.view(self.ep_size, -1)
+ .sum(1)
+ .cpu()
+ .numpy()
+ .tolist()
+ )
+ gathered_tokens = sorted_tokens.new_empty(
+ tokens_per_expert_group.sum(dim=0).cpu().item(), sorted_tokens.shape[1]
+ )
+ input_split_sizes = tokens_per_ep_rank.cpu().numpy().tolist()
+ dist.all_to_all(
+ list(gathered_tokens.split(output_splits)),
+ list(sorted_tokens.split(input_split_sizes)),
+ )
+ tokens_per_expert_post_gather = tokens_per_expert_group.view(
+ self.ep_size, self.experts_per_rank
+ ).sum(dim=0)
+ gatherd_idxs = np.zeros(shape=(gathered_tokens.shape[0],), dtype=np.int32)
+ s = 0
+ for i, k in enumerate(tokens_per_expert_group.cpu().numpy()):
+ gatherd_idxs[s : s + k] = i % self.experts_per_rank
+ s += k
+ gatherd_idxs = gatherd_idxs.argsort()
+ sorted_tokens = gathered_tokens[gatherd_idxs]
+ tokens_per_expert = tokens_per_expert_post_gather
+ tokens_per_expert = tokens_per_expert.cpu().numpy()
+
+ outputs = []
+ start_idx = 0
+ for i, num_tokens in enumerate(tokens_per_expert):
+ end_idx = start_idx + num_tokens
+ if num_tokens == 0:
+ continue
+ expert = self.experts[i + self.ep_rank * self.experts_per_rank]
+ tokens_for_this_expert = sorted_tokens[start_idx:end_idx]
+ expert_out = expert(tokens_for_this_expert)
+ outputs.append(expert_out)
+ start_idx = end_idx
+
+ outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0)
+ if self.ep_size > 1:
+ new_x = torch.empty_like(outs)
+ new_x[gatherd_idxs] = outs
+ gathered_tokens = new_x.new_empty(*sorted_tokens_shape)
+ dist.all_to_all(
+ list(gathered_tokens.split(input_split_sizes)),
+ list(new_x.split(output_splits)),
+ )
+ outs = gathered_tokens
+
+ new_x = torch.empty_like(outs)
+ new_x[idxs] = outs
+ final_out = (
+ new_x.view(*topk_ids.shape, -1)
+ .type(topk_weight.dtype)
+ .mul_(topk_weight.unsqueeze(dim=-1))
+ .sum(dim=1)
+ .type(new_x.dtype)
+ )
+ return final_out
+
+
+# Copied from transformers.models.llama.modeling_llama.repeat_kv
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+ """
+ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+ num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+ """
+ batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+ if n_rep == 1:
+ return hidden_states
+ hidden_states = hidden_states[:, :, None, :, :].expand(
+ batch, num_key_value_heads, n_rep, slen, head_dim
+ )
+ return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+# Copied from transformers.models.llama.modeling_llama.LlamaAttention with Llama->DeepseekV2
+class DeepseekV2Attention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(self, config: DeepseekV2Config, layer_idx: Optional[int] = None):
+ super().__init__()
+ self.config = config
+ self.layer_idx = layer_idx
+ if layer_idx is None:
+ logger.warning_once(
+ f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "
+ "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
+ "when creating this class."
+ )
+
+ self.attention_dropout = config.attention_dropout
+ self.hidden_size = config.hidden_size
+ self.num_heads = config.num_attention_heads
+
+ self.max_position_embeddings = config.max_position_embeddings
+ self.rope_theta = config.rope_theta
+ self.q_lora_rank = config.q_lora_rank
+ self.qk_rope_head_dim = config.qk_rope_head_dim
+ self.kv_lora_rank = config.kv_lora_rank
+ self.v_head_dim = config.v_head_dim
+ self.qk_nope_head_dim = config.qk_nope_head_dim
+ self.q_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim
+
+ self.is_causal = True
+
+ if self.q_lora_rank is None:
+ self.q_proj = nn.Linear(
+ self.hidden_size, self.num_heads * self.q_head_dim, bias=False
+ )
+ else:
+ self.q_a_proj = nn.Linear(
+ self.hidden_size, config.q_lora_rank, bias=config.attention_bias
+ )
+ self.q_a_layernorm = DeepseekV2RMSNorm(config.q_lora_rank)
+ self.q_b_proj = nn.Linear(
+ config.q_lora_rank, self.num_heads * self.q_head_dim, bias=False
+ )
+
+ self.kv_a_proj_with_mqa = nn.Linear(
+ self.hidden_size,
+ config.kv_lora_rank + config.qk_rope_head_dim,
+ bias=config.attention_bias,
+ )
+ self.kv_a_layernorm = DeepseekV2RMSNorm(config.kv_lora_rank)
+ self.kv_b_proj = nn.Linear(
+ config.kv_lora_rank,
+ self.num_heads
+ * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim),
+ bias=False,
+ )
+
+ self.o_proj = nn.Linear(
+ self.num_heads * self.v_head_dim,
+ self.hidden_size,
+ bias=config.attention_bias,
+ )
+ self._init_rope()
+
+ self.softmax_scale = self.q_head_dim ** (-0.5)
+ if self.config.rope_scaling is not None:
+ mscale_all_dim = self.config.rope_scaling.get("mscale_all_dim", 0)
+ scaling_factor = self.config.rope_scaling["factor"]
+ if mscale_all_dim:
+ mscale = yarn_get_mscale(scaling_factor, mscale_all_dim)
+ self.softmax_scale = self.softmax_scale * mscale * mscale
+
+ def _init_rope(self):
+ if self.config.rope_scaling is None:
+ self.rotary_emb = DeepseekV2RotaryEmbedding(
+ self.qk_rope_head_dim,
+ max_position_embeddings=self.max_position_embeddings,
+ base=self.rope_theta,
+ )
+ else:
+ scaling_type = self.config.rope_scaling["type"]
+ scaling_factor = self.config.rope_scaling["factor"]
+ if scaling_type == "linear":
+ self.rotary_emb = DeepseekV2LinearScalingRotaryEmbedding(
+ self.qk_rope_head_dim,
+ max_position_embeddings=self.max_position_embeddings,
+ scaling_factor=scaling_factor,
+ base=self.rope_theta,
+ )
+ elif scaling_type == "dynamic":
+ self.rotary_emb = DeepseekV2DynamicNTKScalingRotaryEmbedding(
+ self.qk_rope_head_dim,
+ max_position_embeddings=self.max_position_embeddings,
+ scaling_factor=scaling_factor,
+ base=self.rope_theta,
+ )
+ elif scaling_type == "yarn":
+ kwargs = {
+ key: self.config.rope_scaling[key]
+ for key in [
+ "original_max_position_embeddings",
+ "beta_fast",
+ "beta_slow",
+ "mscale",
+ "mscale_all_dim",
+ ]
+ if key in self.config.rope_scaling
+ }
+ self.rotary_emb = DeepseekV2YarnRotaryEmbedding(
+ self.qk_rope_head_dim,
+ max_position_embeddings=self.max_position_embeddings,
+ scaling_factor=scaling_factor,
+ base=self.rope_theta,
+ **kwargs,
+ )
+ else:
+ raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
+
+ def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
+ return (
+ tensor.view(bsz, seq_len, self.num_heads, self.v_head_dim)
+ .transpose(1, 2)
+ .contiguous()
+ )
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_value: Optional[Cache] = None,
+ output_attentions: bool = False,
+ use_cache: bool = False,
+ **kwargs,
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+ if "padding_mask" in kwargs:
+ warnings.warn(
+ "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
+ )
+ bsz, q_len, _ = hidden_states.size()
+
+ if self.q_lora_rank is None:
+ q = self.q_proj(hidden_states)
+ else:
+ q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states)))
+ q = q.view(bsz, q_len, self.num_heads, self.q_head_dim).transpose(1, 2)
+ q_nope, q_pe = torch.split(
+ q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1
+ )
+
+ compressed_kv = self.kv_a_proj_with_mqa(hidden_states)
+ compressed_kv, k_pe = torch.split(
+ compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
+ )
+ k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim).transpose(1, 2)
+ kv = (
+ self.kv_b_proj(self.kv_a_layernorm(compressed_kv))
+ .view(bsz, q_len, self.num_heads, self.qk_nope_head_dim + self.v_head_dim)
+ .transpose(1, 2)
+ )
+
+ k_nope, value_states = torch.split(
+ kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1
+ )
+ kv_seq_len = value_states.shape[-2]
+ if past_key_value is not None:
+ if self.layer_idx is None:
+ raise ValueError(
+ f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
+ "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
+ "with a layer index."
+ )
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
+ cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+
+ q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids)
+
+ query_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim)
+ query_states[:, :, :, : self.qk_nope_head_dim] = q_nope
+ query_states[:, :, :, self.qk_nope_head_dim :] = q_pe
+
+ key_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim)
+ key_states[:, :, :, : self.qk_nope_head_dim] = k_nope
+ key_states[:, :, :, self.qk_nope_head_dim :] = k_pe
+ if past_key_value is not None:
+ cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models
+ key_states, value_states = past_key_value.update(
+ key_states, value_states, self.layer_idx, cache_kwargs
+ )
+
+ attn_weights = (
+ torch.matmul(query_states, key_states.transpose(2, 3)) * self.softmax_scale
+ )
+
+ if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
+ raise ValueError(
+ f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
+ f" {attn_weights.size()}"
+ )
+ assert attention_mask is not None
+ if attention_mask is not None:
+ if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
+ raise ValueError(
+ f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
+ )
+ attn_weights = attn_weights + attention_mask
+
+ # upcast attention to fp32
+ attn_weights = nn.functional.softmax(
+ attn_weights, dim=-1, dtype=torch.float32
+ ).to(query_states.dtype)
+ attn_weights = nn.functional.dropout(
+ attn_weights, p=self.attention_dropout, training=self.training
+ )
+ attn_output = torch.matmul(attn_weights, value_states)
+
+ if attn_output.size() != (bsz, self.num_heads, q_len, self.v_head_dim):
+ raise ValueError(
+ f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.v_head_dim)}, but is"
+ f" {attn_output.size()}"
+ )
+
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.v_head_dim)
+
+ attn_output = self.o_proj(attn_output)
+
+ if not output_attentions:
+ attn_weights = None
+
+ return attn_output, attn_weights, past_key_value
+
+
+# Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2 with Llama->DeepseekV2
+class DeepseekV2FlashAttention2(DeepseekV2Attention):
+ """
+ DeepseekV2 flash attention module. This module inherits from `DeepseekV2Attention` as the weights of the module stays
+ untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
+ flash attention and deal with padding tokens in case the input contains any of them.
+ """
+
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+
+ # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
+ # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
+ # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).
+ self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: Optional[torch.LongTensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_value: Optional[Cache] = None,
+ output_attentions: bool = False,
+ use_cache: bool = False,
+ **kwargs,
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+ # DeepseekV2FlashAttention2 attention does not support output_attentions
+ if "padding_mask" in kwargs:
+ warnings.warn(
+ "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
+ )
+
+ # overwrite attention_mask with padding_mask
+ attention_mask = kwargs.pop("padding_mask")
+
+ output_attentions = False
+
+ bsz, q_len, _ = hidden_states.size()
+
+ if self.q_lora_rank is None:
+ # print(f"dtype of hidden_states: {hidden_states.dtype}")
+ # print(f"dtype of q_proj: {self.q_proj.weight.dtype}")
+ q = self.q_proj(hidden_states)
+ else:
+ q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states)))
+ q = q.view(bsz, q_len, self.num_heads, self.q_head_dim).transpose(1, 2)
+ q_nope, q_pe = torch.split(
+ q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1
+ )
+
+ # Flash attention requires the input to have the shape
+ # batch_size x seq_length x head_dim x hidden_dim
+ # therefore we just need to keep the original shape
+ compressed_kv = self.kv_a_proj_with_mqa(hidden_states)
+ compressed_kv, k_pe = torch.split(
+ compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
+ )
+ k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim).transpose(1, 2)
+ kv = (
+ self.kv_b_proj(self.kv_a_layernorm(compressed_kv))
+ .view(bsz, q_len, self.num_heads, self.qk_nope_head_dim + self.v_head_dim)
+ .transpose(1, 2)
+ )
+
+ k_nope, value_states = torch.split(
+ kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1
+ )
+ kv_seq_len = value_states.shape[-2]
+
+ kv_seq_len = value_states.shape[-2]
+ if past_key_value is not None:
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
+
+ cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+ q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids)
+
+ query_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim)
+ query_states[:, :, :, : self.qk_nope_head_dim] = q_nope
+ query_states[:, :, :, self.qk_nope_head_dim :] = q_pe
+
+ key_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim)
+ key_states[:, :, :, : self.qk_nope_head_dim] = k_nope
+ key_states[:, :, :, self.qk_nope_head_dim :] = k_pe
+
+ if self.q_head_dim != self.v_head_dim:
+ value_states = F.pad(value_states, [0, self.q_head_dim - self.v_head_dim])
+
+ if past_key_value is not None:
+ cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models
+ key_states, value_states = past_key_value.update(
+ key_states, value_states, self.layer_idx, cache_kwargs
+ )
+
+ # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache
+ # to be able to avoid many of these transpose/reshape/view.
+ query_states = query_states.transpose(1, 2)
+ key_states = key_states.transpose(1, 2)
+ value_states = value_states.transpose(1, 2)
+
+ dropout_rate = self.attention_dropout if self.training else 0.0
+
+ # In PEFT, usually we cast the layer norms in float32 for training stability reasons
+ # therefore the input hidden states gets silently casted in float32. Hence, we need
+ # cast them back in the correct dtype just to be sure everything works as expected.
+ # This might slowdown training & inference so it is recommended to not cast the LayerNorms
+ # in fp32. (DeepseekV2RMSNorm handles it correctly)
+
+ input_dtype = query_states.dtype
+ if input_dtype == torch.float32:
+ # Handle the case where the model is quantized
+ if hasattr(self.config, "_pre_quantization_dtype"):
+ target_dtype = self.config._pre_quantization_dtype
+ elif torch.is_autocast_enabled():
+ target_dtype = torch.get_autocast_gpu_dtype()
+ else:
+ target_dtype = self.q_proj.weight.dtype if self.q_lora_rank is None else self.q_a_proj.weight.dtype
+
+ logger.warning_once(
+ f"The input hidden states seems to be silently casted in float32, this might be related to"
+ f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
+ f" {target_dtype}."
+ )
+
+ query_states = query_states.to(target_dtype)
+ key_states = key_states.to(target_dtype)
+ value_states = value_states.to(target_dtype)
+
+ attn_output = self._flash_attention_forward(
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ q_len,
+ dropout=dropout_rate,
+ softmax_scale=self.softmax_scale,
+ )
+ if self.q_head_dim != self.v_head_dim:
+ attn_output = attn_output[:, :, :, : self.v_head_dim]
+
+ attn_output = attn_output.reshape(
+ bsz, q_len, self.num_heads * self.v_head_dim
+ ).contiguous()
+ attn_output = self.o_proj(attn_output)
+
+ if not output_attentions:
+ attn_weights = None
+
+ return attn_output, attn_weights, past_key_value
+
+ def _flash_attention_forward(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ query_length,
+ dropout=0.0,
+ softmax_scale=None,
+ ):
+ """
+ Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
+ first unpad the input, then computes the attention scores and pad the final attention scores.
+
+ Args:
+ query_states (`torch.Tensor`):
+ Input query states to be passed to Flash Attention API
+ key_states (`torch.Tensor`):
+ Input key states to be passed to Flash Attention API
+ value_states (`torch.Tensor`):
+ Input value states to be passed to Flash Attention API
+ attention_mask (`torch.Tensor`):
+ The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
+ position of padding tokens and 1 for the position of non-padding tokens.
+ dropout (`int`, *optional*):
+ Attention dropout
+ softmax_scale (`float`, *optional*):
+ The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
+ """
+ if not self._flash_attn_uses_top_left_mask:
+ causal = self.is_causal
+ else:
+ # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in DeepseekV2FlashAttention2 __init__.
+ causal = self.is_causal and query_length != 1
+
+ # Contains at least one padding token in the sequence
+ if attention_mask is not None:
+ batch_size = query_states.shape[0]
+ (
+ query_states,
+ key_states,
+ value_states,
+ indices_q,
+ cu_seq_lens,
+ max_seq_lens,
+ ) = self._upad_input(
+ query_states, key_states, value_states, attention_mask, query_length
+ )
+
+ cu_seqlens_q, cu_seqlens_k = cu_seq_lens
+ max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
+
+ attn_output_unpad = flash_attn_varlen_func(
+ query_states,
+ key_states,
+ value_states,
+ cu_seqlens_q=cu_seqlens_q,
+ cu_seqlens_k=cu_seqlens_k,
+ max_seqlen_q=max_seqlen_in_batch_q,
+ max_seqlen_k=max_seqlen_in_batch_k,
+ dropout_p=dropout,
+ softmax_scale=softmax_scale,
+ causal=causal,
+ )
+
+ attn_output = pad_input(
+ attn_output_unpad, indices_q, batch_size, query_length
+ )
+ else:
+ attn_output = flash_attn_func(
+ query_states,
+ key_states,
+ value_states,
+ dropout,
+ softmax_scale=softmax_scale,
+ causal=causal,
+ )
+
+ return attn_output
+
+ def _upad_input(
+ self, query_layer, key_layer, value_layer, attention_mask, query_length
+ ):
+ indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
+ batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
+
+ key_layer = index_first_axis(
+ key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim),
+ indices_k,
+ )
+ value_layer = index_first_axis(
+ value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim),
+ indices_k,
+ )
+ if query_length == kv_seq_len:
+ query_layer = index_first_axis(
+ query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim),
+ indices_k,
+ )
+ cu_seqlens_q = cu_seqlens_k
+ max_seqlen_in_batch_q = max_seqlen_in_batch_k
+ indices_q = indices_k
+ elif query_length == 1:
+ max_seqlen_in_batch_q = 1
+ cu_seqlens_q = torch.arange(
+ batch_size + 1, dtype=torch.int32, device=query_layer.device
+ ) # There is a memcpy here, that is very bad.
+ indices_q = cu_seqlens_q[:-1]
+ query_layer = query_layer.squeeze(1)
+ else:
+ # The -q_len: slice assumes left padding.
+ attention_mask = attention_mask[:, -query_length:]
+ query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(
+ query_layer, attention_mask
+ )
+
+ return (
+ query_layer,
+ key_layer,
+ value_layer,
+ indices_q,
+ (cu_seqlens_q, cu_seqlens_k),
+ (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
+ )
+
+
+ATTENTION_CLASSES = {
+ "eager": DeepseekV2Attention,
+ "flash_attention_2": DeepseekV2FlashAttention2,
+}
+
+
+class DeepseekV2DecoderLayer(nn.Module):
+ def __init__(self, config: DeepseekV2Config, layer_idx: int):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+
+ self.self_attn = ATTENTION_CLASSES[config._attn_implementation](
+ config=config, layer_idx=layer_idx
+ )
+
+ self.mlp = (
+ DeepseekV2MoE(config)
+ if (
+ config.n_routed_experts is not None
+ and layer_idx >= config.first_k_dense_replace
+ and layer_idx % config.moe_layer_freq == 0
+ )
+ else DeepseekV2MLP(config)
+ )
+ self.input_layernorm = DeepseekV2RMSNorm(
+ config.hidden_size, eps=config.rms_norm_eps
+ )
+ self.post_attention_layernorm = DeepseekV2RMSNorm(
+ config.hidden_size, eps=config.rms_norm_eps
+ )
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
+ output_attentions: Optional[bool] = False,
+ use_cache: Optional[bool] = False,
+ **kwargs,
+ ) -> Tuple[
+ torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]
+ ]:
+ """
+ Args:
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+ attention_mask (`torch.FloatTensor`, *optional*):
+ attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1,
+ query_sequence_length, key_sequence_length)` if default attention is used.
+ output_attentions (`bool`, *optional*):
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+ returned tensors for more detail.
+ use_cache (`bool`, *optional*):
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
+ (see `past_key_values`).
+ past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
+ """
+ if "padding_mask" in kwargs:
+ warnings.warn(
+ "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
+ )
+ residual = hidden_states
+ # print(f"1. dtype of residual: {residual.dtype}")
+
+ hidden_states = self.input_layernorm(hidden_states)
+ # print(f"2. dtype of hidden_states before attn: {hidden_states.dtype}")
+ # Self Attention
+ hidden_states, self_attn_weights, present_key_value = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_value=past_key_value,
+ output_attentions=output_attentions,
+ use_cache=use_cache,
+ **kwargs,
+ )
+ # print(f"3. dtype of hidden_states after attn: {hidden_states.dtype}")
+ hidden_states = residual + hidden_states
+
+ # Fully Connected
+ residual = hidden_states
+ hidden_states = self.post_attention_layernorm(hidden_states)
+ # print(f"4. dtype of hidden_states after post layernorm: {hidden_states.dtype}")
+ hidden_states = self.mlp(hidden_states)
+ # print(f"5. dtype of hidden_states after mlp: {hidden_states.dtype}")
+ hidden_states = residual + hidden_states
+
+ outputs = (hidden_states,)
+
+ if output_attentions:
+ outputs += (self_attn_weights,)
+
+ if use_cache:
+ outputs += (present_key_value,)
+
+ return outputs
+
+
+DeepseekV2_START_DOCSTRING = r"""
+ This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
+ library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
+ etc.)
+
+ This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
+ Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
+ and behavior.
+
+ Parameters:
+ config ([`DeepseekV2Config`]):
+ Model configuration class with all the parameters of the model. Initializing with a config file does not
+ load the weights associated with the model, only the configuration. Check out the
+ [`~PreTrainedModel.from_pretrained`] method to load the model weights.
+"""
+
+
+@add_start_docstrings(
+ "The bare DeepseekV2 Model outputting raw hidden-states without any specific head on top.",
+ DeepseekV2_START_DOCSTRING,
+)
+class DeepseekV2PreTrainedModel(PreTrainedModel):
+ config_class = DeepseekV2Config
+ base_model_prefix = "model"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["DeepseekV2DecoderLayer"]
+ _skip_keys_device_placement = "past_key_values"
+ _supports_flash_attn_2 = True
+ _supports_cache_class = True
+
+ def _init_weights(self, module):
+ std = self.config.initializer_range
+ if isinstance(module, nn.Linear):
+ module.weight.data.normal_(mean=0.0, std=std)
+ if module.bias is not None:
+ module.bias.data.zero_()
+ elif isinstance(module, nn.Embedding):
+ module.weight.data.normal_(mean=0.0, std=std)
+ if module.padding_idx is not None:
+ module.weight.data[module.padding_idx].zero_()
+
+
+DeepseekV2_INPUTS_DOCSTRING = r"""
+ Args:
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+ Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
+ it.
+
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+ [`PreTrainedTokenizer.__call__`] for details.
+
+ [What are input IDs?](../glossary#input-ids)
+ attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+ - 1 for tokens that are **not masked**,
+ - 0 for tokens that are **masked**.
+
+ [What are attention masks?](../glossary#attention-mask)
+
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+ [`PreTrainedTokenizer.__call__`] for details.
+
+ If `past_key_values` is used, optionally only the last `input_ids` have to be input (see
+ `past_key_values`).
+
+ If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
+ and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
+ information on the default strategy.
+
+ - 1 indicates the head is **not masked**,
+ - 0 indicates the head is **masked**.
+ position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
+ config.n_positions - 1]`.
+
+ [What are position IDs?](../glossary#position-ids)
+ past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):
+ Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
+ blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`
+ returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.
+
+ Two formats are allowed:
+ - a [`~cache_utils.Cache`] instance;
+ - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
+ shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy
+ cache format.
+
+ The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the
+ legacy cache format will be returned.
+
+ If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't
+ have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`
+ of shape `(batch_size, sequence_length)`.
+ inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
+ Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
+ is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
+ model's internal embedding lookup matrix.
+ use_cache (`bool`, *optional*):
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
+ `past_key_values`).
+ output_attentions (`bool`, *optional*):
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
+ tensors for more detail.
+ output_hidden_states (`bool`, *optional*):
+ Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
+ more detail.
+ return_dict (`bool`, *optional*):
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+"""
+
+
+@add_start_docstrings(
+ "The bare DeepseekV2 Model outputting raw hidden-states without any specific head on top.",
+ DeepseekV2_START_DOCSTRING,
+)
+class DeepseekV2Model(DeepseekV2PreTrainedModel):
+ """
+ Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`DeepseekV2DecoderLayer`]
+
+ Args:
+ config: DeepseekV2Config
+ """
+
+ def __init__(self, config: DeepseekV2Config):
+ super().__init__(config)
+ self.padding_idx = config.pad_token_id
+ self.vocab_size = config.vocab_size
+
+ self.embed_tokens = nn.Embedding(
+ config.vocab_size, config.hidden_size, self.padding_idx
+ )
+ self.layers = nn.ModuleList(
+ [
+ DeepseekV2DecoderLayer(config, layer_idx)
+ for layer_idx in range(config.num_hidden_layers)
+ ]
+ )
+ self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
+ self.norm = DeepseekV2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+ self.gradient_checkpointing = False
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.embed_tokens
+
+ def set_input_embeddings(self, value):
+ self.embed_tokens = value
+
+ @add_start_docstrings_to_model_forward(DeepseekV2_INPUTS_DOCSTRING)
+ def forward(
+ self,
+ input_ids: torch.LongTensor = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ use_cache: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ ) -> Union[Tuple, BaseModelOutputWithPast]:
+ output_attentions = (
+ output_attentions
+ if output_attentions is not None
+ else self.config.output_attentions
+ )
+ output_hidden_states = (
+ output_hidden_states
+ if output_hidden_states is not None
+ else self.config.output_hidden_states
+ )
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+ return_dict = (
+ return_dict if return_dict is not None else self.config.use_return_dict
+ )
+
+ # retrieve input_ids and inputs_embeds
+ if input_ids is not None and inputs_embeds is not None:
+ raise ValueError(
+ "You cannot specify both input_ids and inputs_embeds at the same time"
+ )
+ elif input_ids is not None:
+ batch_size, seq_length = input_ids.shape[:2]
+ elif inputs_embeds is not None:
+ batch_size, seq_length = inputs_embeds.shape[:2]
+ else:
+ raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+ if self.gradient_checkpointing and self.training:
+ if use_cache:
+ logger.warning_once(
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`transformers."
+ )
+ use_cache = False
+
+ past_key_values_length = 0
+ if use_cache:
+ use_legacy_cache = not isinstance(past_key_values, Cache)
+ if use_legacy_cache:
+ past_key_values = DynamicCache.from_legacy_cache(past_key_values)
+ past_key_values_length = past_key_values.get_usable_length(seq_length)
+
+ if position_ids is None:
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
+ position_ids = torch.arange(
+ past_key_values_length,
+ seq_length + past_key_values_length,
+ dtype=torch.long,
+ device=device,
+ )
+ position_ids = position_ids.unsqueeze(0)
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+
+ if self._use_flash_attention_2:
+ # 2d mask is passed through the layers
+ attention_mask = (
+ attention_mask
+ if (attention_mask is not None and 0 in attention_mask)
+ else None
+ )
+ else:
+ # 4d mask is passed through the layers
+ attention_mask = _prepare_4d_causal_attention_mask(
+ attention_mask,
+ (batch_size, seq_length),
+ inputs_embeds,
+ past_key_values_length,
+ )
+
+ # embed positions
+ hidden_states = inputs_embeds
+
+ # decoder layers
+ all_hidden_states = () if output_hidden_states else None
+ all_self_attns = () if output_attentions else None
+ next_decoder_cache = None
+
+ for decoder_layer in self.layers:
+ if output_hidden_states:
+ all_hidden_states += (hidden_states,)
+
+ if self.gradient_checkpointing and self.training:
+ layer_outputs = self._gradient_checkpointing_func(
+ decoder_layer.__call__,
+ hidden_states,
+ attention_mask,
+ position_ids,
+ past_key_values,
+ output_attentions,
+ use_cache,
+ )
+ else:
+ layer_outputs = decoder_layer(
+ hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_value=past_key_values,
+ output_attentions=output_attentions,
+ use_cache=use_cache,
+ )
+
+ hidden_states = layer_outputs[0]
+
+ if use_cache:
+ next_decoder_cache = layer_outputs[2 if output_attentions else 1]
+
+ if output_attentions:
+ all_self_attns += (layer_outputs[1],)
+
+ hidden_states = self.norm(hidden_states)
+
+ # add hidden states from the last decoder layer
+ if output_hidden_states:
+ all_hidden_states += (hidden_states,)
+
+ next_cache = None
+ if use_cache:
+ next_cache = (
+ next_decoder_cache.to_legacy_cache()
+ if use_legacy_cache
+ else next_decoder_cache
+ )
+ if not return_dict:
+ return tuple(
+ v
+ for v in [hidden_states, next_cache, all_hidden_states, all_self_attns]
+ if v is not None
+ )
+ return BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=next_cache,
+ hidden_states=all_hidden_states,
+ attentions=all_self_attns,
+ )
+
+
+class DeepseekV2ForCausalLM(DeepseekV2PreTrainedModel):
+ _tied_weights_keys = ["lm_head.weight"]
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.model = DeepseekV2Model(config)
+ self.vocab_size = config.vocab_size
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.model.embed_tokens
+
+ def set_input_embeddings(self, value):
+ self.model.embed_tokens = value
+
+ def get_output_embeddings(self):
+ return self.lm_head
+
+ def set_output_embeddings(self, new_embeddings):
+ self.lm_head = new_embeddings
+
+ def set_decoder(self, decoder):
+ self.model = decoder
+
+ def get_decoder(self):
+ return self.model
+
+ @add_start_docstrings_to_model_forward(DeepseekV2_INPUTS_DOCSTRING)
+ @replace_return_docstrings(
+ output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC
+ )
+ def forward(
+ self,
+ input_ids: torch.LongTensor = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ labels: Optional[torch.LongTensor] = None,
+ use_cache: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ ) -> Union[Tuple, CausalLMOutputWithPast]:
+ r"""
+ Args:
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, transformers.,
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+ (masked), the loss is only computed for the tokens with labels in `[0, transformers., config.vocab_size]`.
+
+ Returns:
+
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, DeepseekV2ForCausalLM
+
+ >>> model = DeepseekV2ForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
+ >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
+
+ >>> prompt = "Hey, are you conscious? Can you talk to me?"
+ >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+ >>> # Generate
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+ "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+ ```"""
+ output_attentions = (
+ output_attentions
+ if output_attentions is not None
+ else self.config.output_attentions
+ )
+ output_hidden_states = (
+ output_hidden_states
+ if output_hidden_states is not None
+ else self.config.output_hidden_states
+ )
+ return_dict = (
+ return_dict if return_dict is not None else self.config.use_return_dict
+ )
+
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+ outputs = self.model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ hidden_states = outputs[0]
+ logits = self.lm_head(hidden_states)
+ logits = logits.float()
+
+ loss = None
+ if labels is not None:
+ # Shift so that tokens < n predict n
+ shift_logits = logits[..., :-1, :].contiguous()
+ shift_labels = labels[..., 1:].contiguous()
+ # Flatten the tokens
+ loss_fct = CrossEntropyLoss()
+ shift_logits = shift_logits.view(-1, self.config.vocab_size)
+ shift_labels = shift_labels.view(-1)
+ # Enable model parallelism
+ shift_labels = shift_labels.to(shift_logits.device)
+ loss = loss_fct(shift_logits, shift_labels)
+
+ if not return_dict:
+ output = (logits,) + outputs[1:]
+ return (loss,) + output if loss is not None else output
+
+ return CausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+ def prepare_inputs_for_generation(
+ self,
+ input_ids,
+ past_key_values=None,
+ attention_mask=None,
+ inputs_embeds=None,
+ **kwargs,
+ ):
+ if past_key_values is not None:
+ if isinstance(past_key_values, Cache):
+ cache_length = past_key_values.get_seq_length()
+ past_length = past_key_values.seen_tokens
+ max_cache_length = past_key_values.get_max_length()
+ else:
+ cache_length = past_length = past_key_values[0][0].shape[2]
+ max_cache_length = None
+
+ # Keep only the unprocessed tokens:
+ # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
+ # some of the inputs are exclusivelly passed as part of the cache (e.g. when passing input_embeds as
+ # input)
+ if (
+ attention_mask is not None
+ and attention_mask.shape[1] > input_ids.shape[1]
+ ):
+ input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
+ # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
+ # input_ids based on the past_length.
+ elif past_length < input_ids.shape[1]:
+ input_ids = input_ids[:, past_length:]
+ # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
+
+ # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
+ if (
+ max_cache_length is not None
+ and attention_mask is not None
+ and cache_length + input_ids.shape[1] > max_cache_length
+ ):
+ attention_mask = attention_mask[:, -max_cache_length:]
+
+ position_ids = kwargs.get("position_ids", None)
+ if attention_mask is not None and position_ids is None:
+ # create position_ids on the fly for batch generation
+ position_ids = attention_mask.long().cumsum(-1) - 1
+ position_ids.masked_fill_(attention_mask == 0, 1)
+ if past_key_values:
+ position_ids = position_ids[:, -input_ids.shape[1] :]
+
+ # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+ if inputs_embeds is not None and past_key_values is None:
+ model_inputs = {"inputs_embeds": inputs_embeds}
+ else:
+ model_inputs = {"input_ids": input_ids}
+
+ model_inputs.update(
+ {
+ "position_ids": position_ids,
+ "past_key_values": past_key_values,
+ "use_cache": kwargs.get("use_cache"),
+ "attention_mask": attention_mask,
+ }
+ )
+ return model_inputs
+
+ @staticmethod
+ def _reorder_cache(past_key_values, beam_idx):
+ reordered_past = ()
+ for layer_past in past_key_values:
+ reordered_past += (
+ tuple(
+ past_state.index_select(0, beam_idx.to(past_state.device))
+ for past_state in layer_past
+ ),
+ )
+ return reordered_past
+
+
+@add_start_docstrings(
+ """
+ The DeepseekV2 Model transformer with a sequence classification head on top (linear layer).
+
+ [`DeepseekV2ForSequenceClassification`] uses the last token in order to do the classification, as other causal models
+ (e.g. GPT-2) do.
+
+ Since it does classification on the last token, it requires to know the position of the last token. If a
+ `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
+ no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
+ padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
+ each row of the batch).
+ """,
+ DeepseekV2_START_DOCSTRING,
+)
+class DeepseekV2ForSequenceClassification(DeepseekV2PreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+ self.model = DeepseekV2Model(config)
+ self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.model.embed_tokens
+
+ def set_input_embeddings(self, value):
+ self.model.embed_tokens = value
+
+ @add_start_docstrings_to_model_forward(DeepseekV2_INPUTS_DOCSTRING)
+ def forward(
+ self,
+ input_ids: torch.LongTensor = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ labels: Optional[torch.LongTensor] = None,
+ use_cache: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, transformers.,
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+ """
+ return_dict = (
+ return_dict if return_dict is not None else self.config.use_return_dict
+ )
+
+ transformer_outputs = self.model(
+ input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+ hidden_states = transformer_outputs[0]
+ logits = self.score(hidden_states)
+
+ if input_ids is not None:
+ batch_size = input_ids.shape[0]
+ else:
+ batch_size = inputs_embeds.shape[0]
+
+ if self.config.pad_token_id is None and batch_size != 1:
+ raise ValueError(
+ "Cannot handle batch sizes > 1 if no padding token is defined."
+ )
+ if self.config.pad_token_id is None:
+ sequence_lengths = -1
+ else:
+ if input_ids is not None:
+ sequence_lengths = (
+ torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
+ ).to(logits.device)
+ else:
+ sequence_lengths = -1
+
+ pooled_logits = logits[
+ torch.arange(batch_size, device=logits.device), sequence_lengths
+ ]
+
+ loss = None
+ if labels is not None:
+ labels = labels.to(logits.device)
+ if self.config.problem_type is None:
+ if self.num_labels == 1:
+ self.config.problem_type = "regression"
+ elif self.num_labels > 1 and (
+ labels.dtype == torch.long or labels.dtype == torch.int
+ ):
+ self.config.problem_type = "single_label_classification"
+ else:
+ self.config.problem_type = "multi_label_classification"
+
+ if self.config.problem_type == "regression":
+ loss_fct = MSELoss()
+ if self.num_labels == 1:
+ loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
+ else:
+ loss = loss_fct(pooled_logits, labels)
+ elif self.config.problem_type == "single_label_classification":
+ loss_fct = CrossEntropyLoss()
+ loss = loss_fct(
+ pooled_logits.view(-1, self.num_labels), labels.view(-1)
+ )
+ elif self.config.problem_type == "multi_label_classification":
+ loss_fct = BCEWithLogitsLoss()
+ loss = loss_fct(pooled_logits, labels)
+ if not return_dict:
+ output = (pooled_logits,) + transformer_outputs[1:]
+ return ((loss,) + output) if loss is not None else output
+
+ return SequenceClassifierOutputWithPast(
+ loss=loss,
+ logits=pooled_logits,
+ past_key_values=transformer_outputs.past_key_values,
+ hidden_states=transformer_outputs.hidden_states,
+ attentions=transformer_outputs.attentions,
+ )
diff --git a/mftcoder_accelerate/src/model/deepseek_v2/tokenization_deepseek_fast.py b/mftcoder_accelerate/src/model/deepseek_v2/tokenization_deepseek_fast.py
new file mode 100644
index 0000000..d243771
--- /dev/null
+++ b/mftcoder_accelerate/src/model/deepseek_v2/tokenization_deepseek_fast.py
@@ -0,0 +1,38 @@
+from typing import List, Optional, Union
+
+
+from transformers.models.llama import LlamaTokenizerFast
+
+
+class DeepseekTokenizerFast(LlamaTokenizerFast):
+
+ def convert_ids_to_tokens(
+ self, ids: Union[int, List[int]], skip_special_tokens: bool = False
+ ) -> Union[str, List[str]]:
+ """
+ Converts a single index or a sequence of indices in a token or a sequence of tokens, using the vocabulary and
+ added tokens.
+
+ Args:
+ ids (`int` or `List[int]`):
+ The token id (or token ids) to convert to tokens.
+ skip_special_tokens (`bool`, *optional*, defaults to `False`):
+ Whether or not to remove special tokens in the decoding.
+
+ Returns:
+ `str` or `List[str]`: The decoded token(s).
+ """
+ if isinstance(ids, int):
+ return self._convert_id_to_token(ids)
+ tokens = []
+ for index in ids:
+ index = int(index)
+ if skip_special_tokens and index in self.all_special_ids:
+ continue
+ token = self._tokenizer.id_to_token(index)
+ tokens.append(token if token is not None else "")
+ return tokens
+
+ def _convert_id_to_token(self, index: int) -> Optional[str]:
+ token = self._tokenizer.id_to_token(int(index))
+ return token if token is not None else ""
diff --git a/mft_peft_hf/src/model/gpt_bigcode/__init__.py b/mftcoder_accelerate/src/model/gpt_bigcode/__init__.py
similarity index 100%
rename from mft_peft_hf/src/model/gpt_bigcode/__init__.py
rename to mftcoder_accelerate/src/model/gpt_bigcode/__init__.py
diff --git a/mft_peft_hf/src/model/gpt_bigcode/configuration_gpt_bigcode.py b/mftcoder_accelerate/src/model/gpt_bigcode/configuration_gpt_bigcode.py
similarity index 100%
rename from mft_peft_hf/src/model/gpt_bigcode/configuration_gpt_bigcode.py
rename to mftcoder_accelerate/src/model/gpt_bigcode/configuration_gpt_bigcode.py
diff --git a/mft_peft_hf/src/model/gpt_bigcode/modeling_gpt_bigcode.py b/mftcoder_accelerate/src/model/gpt_bigcode/modeling_gpt_bigcode.py
similarity index 100%
rename from mft_peft_hf/src/model/gpt_bigcode/modeling_gpt_bigcode.py
rename to mftcoder_accelerate/src/model/gpt_bigcode/modeling_gpt_bigcode.py
diff --git a/mft_atorch/model/gpt_neox/__init__.py b/mftcoder_accelerate/src/model/gpt_neox/__init__.py
similarity index 100%
rename from mft_atorch/model/gpt_neox/__init__.py
rename to mftcoder_accelerate/src/model/gpt_neox/__init__.py
diff --git a/mft_peft_hf/src/model/gpt_neox/config.json b/mftcoder_accelerate/src/model/gpt_neox/config.json
similarity index 100%
rename from mft_peft_hf/src/model/gpt_neox/config.json
rename to mftcoder_accelerate/src/model/gpt_neox/config.json
diff --git a/mft_peft_hf/src/model/gpt_neox/configuration_gpt_neox.py b/mftcoder_accelerate/src/model/gpt_neox/configuration_gpt_neox.py
similarity index 100%
rename from mft_peft_hf/src/model/gpt_neox/configuration_gpt_neox.py
rename to mftcoder_accelerate/src/model/gpt_neox/configuration_gpt_neox.py
diff --git a/mft_atorch/model/gpt_neox/generation_config.json b/mftcoder_accelerate/src/model/gpt_neox/generation_config.json
similarity index 100%
rename from mft_atorch/model/gpt_neox/generation_config.json
rename to mftcoder_accelerate/src/model/gpt_neox/generation_config.json
diff --git a/mft_peft_hf/src/model/gpt_neox/modeling_gpt_neox.py b/mftcoder_accelerate/src/model/gpt_neox/modeling_gpt_neox.py
similarity index 100%
rename from mft_peft_hf/src/model/gpt_neox/modeling_gpt_neox.py
rename to mftcoder_accelerate/src/model/gpt_neox/modeling_gpt_neox.py
diff --git a/mft_peft_hf/src/model/gpt_neox/tokenization_gpt_neox_fast.py b/mftcoder_accelerate/src/model/gpt_neox/tokenization_gpt_neox_fast.py
similarity index 100%
rename from mft_peft_hf/src/model/gpt_neox/tokenization_gpt_neox_fast.py
rename to mftcoder_accelerate/src/model/gpt_neox/tokenization_gpt_neox_fast.py
diff --git a/mftcoder_accelerate/src/model/phi/configuration_mixformer_sequential.py b/mftcoder_accelerate/src/model/phi/configuration_mixformer_sequential.py
new file mode 100644
index 0000000..8cc2d51
--- /dev/null
+++ b/mftcoder_accelerate/src/model/phi/configuration_mixformer_sequential.py
@@ -0,0 +1,53 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+import math
+from typing import Any, Dict, List, Optional, Union
+
+from transformers import PretrainedConfig
+
+
+class MixFormerSequentialConfig(PretrainedConfig):
+ """MixFormer (sequential for DeepSpeed) configuration."""
+
+ model_type = "mixformer-sequential"
+
+ attribute_map = {
+ "max_position_embeddings": "n_positions",
+ "hidden_size": "n_embd",
+ "num_attention_heads": "n_head",
+ "num_hidden_layers": "n_layer",
+ }
+
+ def __init__(
+ self,
+ vocab_size: Optional[int] = 50304,
+ n_positions: Optional[int] = 2048,
+ n_embd: Optional[int] = 1024,
+ n_layer: Optional[int] = 20,
+ n_inner: Optional[int] = None,
+ n_head: Optional[int] = 16,
+ rotary_dim: Optional[int] = 32,
+ activation_function: Optional[str] = "gelu_new",
+ embd_pdrop: Optional[float] = 0.0,
+ resid_pdrop: Optional[float] = 0.0,
+ layer_norm_epsilon: Optional[float] = 1e-5,
+ initializer_range: Optional[float] = 0.02,
+ tie_word_embeddings: Optional[bool] = False,
+ pad_vocab_size_multiple: Optional[int] = 64,
+ **kwargs
+ ) -> None:
+ self.vocab_size = int(math.ceil(vocab_size / pad_vocab_size_multiple) * pad_vocab_size_multiple)
+ self.n_positions = n_positions
+ self.n_embd = n_embd
+ self.n_layer = n_layer
+ self.n_inner = n_inner
+ self.n_head = n_head
+ self.rotary_dim = min(rotary_dim, n_embd // n_head)
+ self.activation_function = activation_function
+ self.embd_pdrop = embd_pdrop
+ self.resid_pdrop = resid_pdrop
+ self.layer_norm_epsilon = layer_norm_epsilon
+ self.initializer_range = initializer_range
+
+ super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
diff --git a/mftcoder_accelerate/src/model/phi/modeling_mixformer_sequential.py b/mftcoder_accelerate/src/model/phi/modeling_mixformer_sequential.py
new file mode 100644
index 0000000..4b26a55
--- /dev/null
+++ b/mftcoder_accelerate/src/model/phi/modeling_mixformer_sequential.py
@@ -0,0 +1,764 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+#
+# BSD 3-Clause License
+#
+# Copyright (c) 2022, Tri Dao, trid@cs.stanford.edu.
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# * Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# * Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# * Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+from __future__ import annotations
+
+import math
+import copy
+from typing import Any, Dict, Optional, Tuple, Union
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn as nn
+
+from einops import rearrange
+from transformers.activations import ACT2FN
+from transformers import PretrainedConfig, PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from .configuration_mixformer_sequential import MixFormerSequentialConfig
+
+import xformers.ops
+
+@dataclass
+class InferenceParams:
+ """Inference parameters passed to model to efficiently calculate
+ and store context during inference.
+ Reference:
+ https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/utils/generation.py.
+ Args:
+ max_sequence_len: Maximum sequence length.
+ max_batch_size: Maximum batch size.
+ sequence_len_offset: Sequence length offset.
+ batch_size_offset: Batch size offset.
+ key_value_memory_dict: Key value memory dictionary.
+ fused_ft_kernel: Whether to use fused kernel for fast inference.
+ lengths_per_sample: Lengths per sample.
+ """
+
+ max_sequence_len: int = field(metadata={"help": "Maximum sequence length."})
+
+ max_batch_size: int = field(metadata={"help": "Maximum batch size."})
+
+ sequence_len_offset: int = field(default=0, metadata={"help": "Sequence length offset."})
+
+ batch_size_offset: int = field(default=0, metadata={"help": "Batch size offset."})
+
+ key_value_memory_dict: Dict[str, Any] = field(
+ default_factory=dict, metadata={"help": "Key value memory dictionary."}
+ )
+
+ fused_ft_kernel: bool = field(default=False, metadata={"help": "Whether to use fused kernel for fast inference."})
+
+ lengths_per_sample: torch.Tensor = field(default=None, metadata={"help": "Lengths per sample."})
+
+
+class Embedding(nn.Module):
+ """Token embedding with dropout."""
+
+ def __init__(self, config: PretrainedConfig) -> None:
+ super().__init__()
+
+ self.wte = nn.Embedding(config.vocab_size, config.n_embd)
+ self.drop = nn.Dropout(config.embd_pdrop)
+
+ def forward(self, input_ids: torch.LongTensor) -> torch.FloatTensor:
+ input_shape = input_ids.size()
+ input_ids = input_ids.view(-1, input_shape[-1])
+
+ hidden_states = self.wte(input_ids)
+ hidden_states = self.drop(hidden_states)
+
+ return hidden_states
+
+
+class RotaryEmbedding(nn.Module):
+ """Rotary embeddings.
+ Reference:
+ https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/layers/rotary.py.
+
+ """
+
+ def __init__(
+ self,
+ dim: int,
+ base: int = 10000,
+ scale_base: Optional[float] = None,
+ device: Optional[str] = None,
+ **kwargs,
+ ) -> None:
+ super().__init__()
+
+ if scale_base is not None:
+ raise NotImplementedError
+
+ # Generate and save the inverse frequency buffer (non-trainable)
+ self.dim = dim
+ self.base = base
+ self.scale_base = scale_base
+ self.device = device
+
+ inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, device=device, dtype=torch.float32) / dim))
+ self.register_buffer("inv_freq", inv_freq)
+
+ scale = (
+ (torch.arange(0, dim, 2, device=device, dtype=torch.float32) + 0.4 * dim) / (1.4 * dim)
+ if scale_base is not None
+ else None
+ )
+ self.register_buffer("scale", scale)
+
+ self._seq_len_cached = 0
+ self._cos_cached = None
+ self._sin_cached = None
+ self._cos_k_cached = None
+ self._sin_k_cached = None
+
+ def _update_cos_sin_cache(self, x: torch.FloatTensor, seqlen_offset: int = 0) -> None:
+ # Reset the tables if the sequence length has changed,
+ # or if we're on a new device (possibly due to tracing for instance)
+ seqlen = x.shape[1] + seqlen_offset
+
+ # Re-generate the inverse frequency buffer if it's not fp32
+ # (for instance if model.half() was called)
+ if self.inv_freq.dtype != "torch.float32":
+ self.inv_freq = 1.0 / (
+ self.base ** (torch.arange(0, self.dim, 2, device=self.device, dtype=torch.float32) / self.dim)
+ )
+
+ if seqlen > self._seq_len_cached or self._cos_cached.device != x.device or self._cos_cached.dtype != x.dtype:
+ self._seq_len_cached = seqlen
+ t = torch.arange(seqlen, device=x.device, dtype=torch.float32)
+
+ # Don't do einsum, it converts fp32 to fp16
+ # freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+ freqs = torch.outer(t, self.inv_freq.to(device=t.device, dtype=torch.float32))
+ if self.scale is None:
+ self._cos_cached = torch.cos(freqs).to(x.dtype)
+ self._sin_cached = torch.sin(freqs).to(x.dtype)
+ else:
+ power = (
+ torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - seqlen // 2
+ ) / self.scale_base
+ scale = self.scale.to(device=power.device) ** rearrange(power, "s -> s 1")
+
+ # We want the multiplication by scale to happen in fp32
+ self._cos_cached = (torch.cos(freqs) * scale).to(x.dtype)
+ self._sin_cached = (torch.sin(freqs) * scale).to(x.dtype)
+ self._cos_k_cached = (torch.cos(freqs) / scale).to(x.dtype)
+ self._sin_k_cached = (torch.sin(freqs) / scale).to(x.dtype)
+
+ def _apply_rotary_emb_qkv(
+ self,
+ qkv: torch.FloatTensor,
+ sin: torch.FloatTensor,
+ cos: torch.FloatTensor,
+ sin_k: Optional[torch.FloatTensor] = None,
+ cos_k: Optional[torch.FloatTensor] = None,
+ ) -> torch.FloatTensor:
+ _, seqlen, three, _, headdim = qkv.shape
+ assert three == 3
+
+ rotary_seqlen, rotary_dim = cos.shape
+ rotary_dim *= 2
+ assert rotary_dim <= headdim
+ assert seqlen <= rotary_seqlen
+
+ cos_k = cos if cos_k is None else cos_k
+ sin_k = sin if sin_k is None else sin_k
+ assert sin.shape == cos_k.shape == sin_k.shape == (rotary_seqlen, rotary_dim // 2)
+
+ q_rot = qkv[:, :, 0, :, :rotary_dim]
+ q_pass = qkv[:, :, 0, :, rotary_dim:]
+
+ k_rot = qkv[:, :, 1, :, :rotary_dim]
+ k_pass = qkv[:, :, 1, :, rotary_dim:]
+
+ # Splits the queries and keys in half
+ q1, q2 = q_rot.chunk(2, dim=-1)
+ k1, k2 = k_rot.chunk(2, dim=-1)
+ c, s = rearrange(cos[:seqlen], "s d -> s 1 d"), rearrange(sin[:seqlen], "s d -> s 1 d")
+
+ # Casts to fp32 are necessary to prevent fp16 overflow issues
+ q1, q2, k1, k2, c, s = [t.to(dtype=torch.float32) for t in [q1, q2, k1, k2, c, s]]
+
+ # Computes the new keys and queries, recasting to original dtype
+ q_rot = torch.cat([q1 * c - q2 * s, q1 * s + q2 * c], axis=-1).to(qkv.dtype)
+ k_rot = torch.cat([k1 * c - k2 * s, k1 * s + k2 * c], axis=-1).to(qkv.dtype)
+
+ return torch.cat(
+ [
+ torch.cat([q_rot, q_pass], axis=-1).unsqueeze(2),
+ torch.cat([k_rot, k_pass], axis=-1).unsqueeze(2),
+ qkv[:, :, 2:3, :, :],
+ ],
+ axis=2,
+ )
+
+ def forward(self, qkv: torch.Tensor, seqlen_offset: int = 0) -> Tuple[torch.Tensor, torch.Tensor]:
+ # `qkv` is of shape (batch, seqlen, 3, nheads, headdim)
+ self._update_cos_sin_cache(qkv, seqlen_offset)
+ return self._apply_rotary_emb_qkv(qkv, self._sin_cached[seqlen_offset:], self._cos_cached[seqlen_offset:])
+
+
+class MLP(nn.Module):
+ """Multi-Layer Perceptron.
+ Reference:
+ Attention Is All You Need.
+ https://arxiv.org/pdf/1706.03762.pdf.
+ """
+
+ def __init__(self, config: PretrainedConfig, n_inner: Optional[int] = None, act_fn: Optional[str] = None) -> None:
+ super().__init__()
+
+ act_fn = config.activation_function if act_fn is None else act_fn
+ assert act_fn in ACT2FN.keys(), f"`act_fn` must be one of: {ACT2FN.keys()}."
+
+ n_inner = getattr(config, "n_inner", None) if n_inner is None else n_inner
+ n_inner = n_inner if n_inner is not None else 4 * config.n_embd
+
+ self.fc1 = nn.Linear(config.n_embd, n_inner)
+ self.fc2 = nn.Linear(n_inner, config.n_embd)
+ self.act = ACT2FN[act_fn]
+
+ def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:
+ hidden_states = self.fc1(hidden_states)
+ hidden_states = self.act(hidden_states)
+ hidden_states = self.fc2(hidden_states)
+
+ return hidden_states
+
+
+class SelfAttention(nn.Module):
+ """Self-attention layer (compatible with PyTorch).
+
+ Reference:
+ https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/modules/mha.py.
+ """
+
+ def __init__(
+ self,
+ causal: bool = True,
+ softmax_scale: Optional[float] = None,
+ attention_dropout: float = 0.0,
+ ) -> None:
+ super().__init__()
+
+ self.causal = causal
+ self.softmax_scale = softmax_scale
+ self.drop = nn.Dropout(attention_dropout)
+
+ def forward(
+ self,
+ qkv: torch.FloatTensor,
+ causal: bool = None,
+ attention_mask: Optional[torch.BoolTensor] = None,
+ **kwargs,
+ ) -> torch.FloatTensor:
+ causal = self.causal if causal is None else causal
+ batch_size, seq_len = qkv.shape[0], qkv.shape[1]
+ q, k, v = qkv.unbind(dim=2)
+
+ # flash attention
+ output = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=xformers.ops.LowerTriangularMask(),
+ op=xformers.ops.MemoryEfficientAttentionFlashAttentionOp)
+
+ # softmax_scale = self.softmax_scale or 1.0 / math.sqrt(q.shape[-1])
+ # scores = torch.einsum("bthd,bshd->bhts", q, k * softmax_scale)
+
+ # if attention_mask is not None:
+ # padding_mask = torch.full((batch_size, seq_len), -10000.0, dtype=scores.dtype, device=scores.device)
+ # padding_mask.masked_fill_(attention_mask, 0.0)
+
+ # scores = scores + rearrange(padding_mask, "b s -> b 1 1 s")
+
+ # if causal:
+ # causal_mask = torch.triu(torch.full((seq_len, seq_len), -10000.0, device=scores.device), 1)
+ # scores = scores + causal_mask.to(dtype=scores.dtype)
+
+ # attention = torch.softmax(scores, dim=-1, dtype=v.dtype)
+ # attention = self.drop(attention)
+
+ # output = torch.einsum("bhts,bshd->bthd", attention, v)
+
+ return output
+
+
+class CrossAttention(nn.Module):
+ """Cross-attention layer (compatible with PyTorch).
+
+ Reference:
+ https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/modules/mha.py.
+
+ """
+
+ def __init__(
+ self,
+ causal: bool = True,
+ softmax_scale: Optional[float] = None,
+ attention_dropout: float = 0.0,
+ ) -> None:
+ super().__init__()
+
+ self.causal = causal
+ self.softmax_scale = softmax_scale
+ self.drop = nn.Dropout(attention_dropout)
+
+ def forward(
+ self,
+ q: torch.FloatTensor,
+ kv: torch.FloatTensor,
+ causal: bool = None,
+ attention_mask: Optional[torch.BoolTensor] = None,
+ **kwargs,
+ ) -> torch.FloatTensor:
+ causal = self.causal if causal is None else causal
+ batch_size, seq_len_q = q.shape[0], q.shape[1]
+ assert kv.shape[0] == batch_size and kv.shape[3] == q.shape[2] and kv.shape[4] == q.shape[3]
+
+ seq_len_k = kv.shape[1]
+ k, v = kv.unbind(dim=2)
+
+ softmax_scale = self.softmax_scale or 1.0 / math.sqrt(q.shape[-1])
+ scores = torch.einsum("bthd,bshd->bhts", q, k * softmax_scale)
+
+ if attention_mask is not None:
+ padding_mask = torch.full((batch_size, seq_len_k), -10000.0, dtype=scores.dtype, device=scores.device)
+ padding_mask.masked_fill_(attention_mask, 0.0)
+
+ scores = scores + rearrange(padding_mask, "b s -> b 1 1 s")
+
+ if causal:
+ causal_mask = torch.triu(torch.full((seq_len_q, seq_len_k), -10000.0, device=scores.device), 1)
+ scores = scores + causal_mask.to(dtype=scores.dtype)
+
+ attention = torch.softmax(scores, dim=-1, dtype=v.dtype)
+ attention = self.drop(attention)
+
+ output = torch.einsum("bhts,bshd->bthd", attention, v)
+
+ return output
+
+
+def find_mha_dims(
+ config: PretrainedConfig, n_head: Optional[int] = None, head_dim: Optional[int] = None
+) -> Tuple[int, int]:
+ """Validate and return the number of heads and head dimension for multi-head attention.
+ Args:
+ config: Model configuration.
+ n_head: Number of heads.
+ head_dim: Head dimension.
+ Returns:
+ Number of heads and head dimension.
+ """
+
+ assert all(
+ hasattr(config, attr) for attr in ["n_embd", "n_head"]
+ ), "`config` must have `n_embd` and `n_head` attributes."
+
+ if head_dim is None:
+ assert (
+ config.n_embd % config.n_head == 0
+ ), f"Hidden size ({config.n_embd}) must be divisible by the number of heads ({config.n_head})."
+
+ if n_head is None and head_dim is None:
+ head_dim = config.n_embd // config.n_head
+ n_head = config.n_head
+ elif n_head is None or head_dim is None:
+ raise ValueError("`n_head` and `head_dim` must be both specified or `None`.")
+
+ return n_head, head_dim
+
+
+def update_kv_cache(kv: torch.FloatTensor, inference_params: InferenceParams, layer_idx: int) -> torch.FloatTensor:
+ """Update the key-value cache for inference.
+ Reference:
+ https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/modules/mha.py.
+ Args:
+ kv: Key-value tensor.
+ inference_params: Inference parameters.
+ layer_idx: Layer index.
+ Returns:
+ Updated key-value tensor.
+ """
+
+ num_heads, head_dim = kv.shape[-2:]
+
+ if layer_idx not in inference_params.key_value_memory_dict:
+ kv_cache = torch.empty(
+ inference_params.max_batch_size,
+ inference_params.max_sequence_len,
+ 2,
+ num_heads,
+ head_dim,
+ dtype=kv.dtype,
+ device=kv.device,
+ )
+ inference_params.key_value_memory_dict[layer_idx] = kv_cache
+ else:
+ if not inference_params.fused_ft_kernel:
+ kv_cache = inference_params.key_value_memory_dict[layer_idx]
+ else:
+ k_cache, v_cache = inference_params.key_value_memory_dict[layer_idx]
+ kv_cache = None
+
+ batch_start = inference_params.batch_size_offset
+ batch_end = batch_start + kv.shape[0]
+ assert batch_end <= (kv_cache.shape[0] if kv_cache is not None else v_cache.shape[0])
+
+ sequence_start = inference_params.sequence_len_offset
+ sequence_end = sequence_start + kv.shape[1]
+ assert sequence_end <= (kv_cache.shape[1] if kv_cache is not None else v_cache.shape[2])
+
+ if not inference_params.fused_ft_kernel:
+ assert kv_cache is not None
+
+ kv_cache[batch_start:batch_end, sequence_start:sequence_end, ...] = kv
+ kv = kv_cache[batch_start:batch_end, :sequence_end, ...]
+
+ return kv
+
+ assert inference_params.sequence_len_offset == 0
+ assert kv.dtype in [torch.float16, torch.bfloat16, torch.float32]
+
+ packsize = 4 if kv.dtype == torch.float32 else 8
+
+ if kv_cache is not None:
+ kv_cache[batch_start:batch_end, sequence_start:sequence_end, ...] = kv
+ k_cache = rearrange(kv_cache[:, :, 0], "b s h (d packsize) -> b h d s packsize", packsize=packsize).contiguous()
+ v_cache = rearrange(kv_cache[:, :, 1], "b s h d -> b h s d").contiguous()
+ inference_params.key_value_memory_dict[layer_idx] = (k_cache, v_cache)
+ else:
+ k_cache[batch_start:batch_end, :, :, :sequence_end, :] = rearrange(
+ kv[:, :, 0], "b s h (d packsize) -> b h d s packsize", packsize=packsize
+ )
+ v_cache[batch_start:batch_end, :, :sequence_end, :] = rearrange(kv[:, :, 1], "b s h d -> b h s d")
+
+ return kv
+
+
+class MHA(nn.Module):
+ """Multi-head attention layer."""
+
+ def __init__(
+ self,
+ config: PretrainedConfig,
+ dtype: Optional[torch.dtype] = None,
+ device: Optional[str] = None,
+ rotary_dim: Optional[int] = None,
+ rotary_emb_scale_base: Optional[float] = None,
+ n_head: Optional[int] = None,
+ head_dim: Optional[int] = None,
+ bias: bool = True,
+ causal: bool = True,
+ softmax_scale: Optional[float] = None,
+ dropout: float = 0.0,
+ layer_idx: Optional[int] = None,
+ return_residual: bool = False,
+ checkpointing: bool = False,
+ ) -> None:
+ super().__init__()
+
+ # Rotary embedding
+ self.rotary_emb_dim = rotary_dim if rotary_dim is not None else getattr(config, "rotary_dim", 0)
+ if self.rotary_emb_dim > 0:
+ rotary_kwargs = {"device": device}
+ if rotary_emb_scale_base is not None and rotary_emb_scale_base > 0.0:
+ rotary_kwargs["scale_base"] = rotary_emb_scale_base
+ self.rotary_emb = RotaryEmbedding(self.rotary_emb_dim, **rotary_kwargs)
+
+ # MLP
+ self.n_head, self.head_dim = find_mha_dims(config, n_head, head_dim)
+ op_size = self.n_head * self.head_dim
+ hidden_size = config.n_embd
+
+ self.Wqkv = nn.Linear(hidden_size, 3 * op_size, bias=bias, device=device, dtype=dtype)
+ self.out_proj = nn.Linear(op_size, hidden_size, bias=bias, device=device, dtype=dtype)
+
+ # Attention
+ self.inner_attn = SelfAttention(causal=causal, softmax_scale=softmax_scale, attention_dropout=dropout)
+ self.inner_cross_attn = CrossAttention(causal=causal, softmax_scale=softmax_scale, attention_dropout=dropout)
+
+ self.layer_idx = layer_idx
+ self.return_residual = return_residual
+ # self.checkpointing = checkpointing
+ self.checkpointing = True
+
+ def forward(
+ self,
+ x: torch.FloatTensor,
+ past_key_values: Optional[InferenceParams] = None,
+ attention_mask: Optional[torch.BoolTensor] = None,
+ cu_seqlens: Optional[torch.LongTensor] = None,
+ max_seqlen: Optional[int] = None,
+ **kwargs,
+ ) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
+ qkv = self.Wqkv(x)
+ qkv = rearrange(qkv, "... (three h d) -> ... three h d", three=3, d=self.head_dim)
+
+ seqlen_offset = past_key_values.sequence_len_offset if past_key_values is not None else 0
+ if self.rotary_emb_dim > 0:
+ qkv = self.rotary_emb(qkv, seqlen_offset=seqlen_offset)
+
+ if past_key_values is not None:
+ kv = update_kv_cache(qkv[:, :, 1:], past_key_values, self.layer_idx)
+
+ if attention_mask is not None:
+ attention_mask = attention_mask[0] if isinstance(attention_mask, tuple) else attention_mask
+ attention_mask = attention_mask.bool().to(qkv.device)
+
+ attention_kwargs = {"attention_mask": attention_mask}
+
+ if past_key_values is None or seqlen_offset == 0:
+ if self.checkpointing:
+ # attn_output = torch.utils.checkpoint.checkpoint(self.inner_attn, qkv, **attention_kwargs)
+ attn_output = torch.utils.checkpoint.checkpoint(self.inner_attn, qkv)
+ else:
+ attn_output = self.inner_attn(qkv, **attention_kwargs)
+ else:
+ q = qkv[:, :, 0]
+ causal = None if past_key_values.sequence_len_offset == 0 else False
+ attn_output = self.inner_cross_attn(q, kv, causal=causal, **attention_kwargs)
+
+ output = rearrange(attn_output, "... h d -> ... (h d)")
+ output = self.out_proj(output)
+
+ return output if not self.return_residual else (output, x)
+
+
+class ParallelBlock(nn.Module):
+ """Parallel block.
+ This block applies parallel mixer and MLP layers to the input (used in GPT-J and CodeGen).
+ """
+
+ def __init__(
+ self,
+ config: PretrainedConfig,
+ block_idx: Optional[int] = None,
+ checkpointing: bool = False,
+ ) -> None:
+ super().__init__()
+
+ self.ln = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
+ self.resid_dropout = nn.Dropout(config.resid_pdrop)
+ self.block_idx = block_idx
+
+ self.mixer = MHA(config, layer_idx=block_idx, checkpointing=checkpointing)
+ self.mlp = MLP(config)
+
+ def forward(
+ self,
+ hidden_states: torch.FloatTensor,
+ past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,
+ attention_mask: Optional[torch.BoolTensor] = None,
+ **kwargs,
+ ) -> torch.FloatTensor:
+ residual = hidden_states
+ hidden_states = self.ln(hidden_states)
+
+ attn_outputs = self.mixer(hidden_states, past_key_values=past_key_values, attention_mask=attention_mask)
+ if isinstance(attn_outputs, tuple):
+ attn_outputs = attn_outputs[0]
+
+ attn_outputs = self.resid_dropout(attn_outputs)
+ feed_forward_hidden_states = self.resid_dropout(self.mlp(hidden_states))
+
+ hidden_states = attn_outputs + feed_forward_hidden_states + residual
+
+ return hidden_states
+
+
+class CausalLMHead(nn.Module):
+ """Causal Language Modeling head.
+ Reference:
+ Improving Language Understanding by Generative Pre-Training.
+ https://cdn.openai.com/research-covers/language-unsupervised/language_understanding_paper.pdf.
+ """
+
+ def __init__(self, config: PretrainedConfig) -> None:
+ super().__init__()
+
+ self.ln = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
+ self.linear = nn.Linear(config.n_embd, config.vocab_size)
+
+ def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:
+ hidden_states = self.ln(hidden_states)
+ logits = self.linear(hidden_states).to(torch.float32)
+
+ return logits
+
+
+class CausalLMLoss(nn.Module):
+ """Causal Language Modeling loss.
+ Reference:
+ Improving Language Understanding by Generative Pre-Training.
+ https://cdn.openai.com/research-covers/language-unsupervised/language_understanding_paper.pdf.
+ """
+
+ def __init__(self, shift_labels: bool = True) -> None:
+ super().__init__()
+
+ self.shift_labels = shift_labels
+ self.loss_fct = nn.CrossEntropyLoss()
+
+ def forward(self, logits: torch.FloatTensor, labels: torch.LongTensor) -> torch.FloatTensor:
+ if self.shift_labels:
+ logits = logits[..., :-1, :].contiguous()
+ labels = labels[..., 1:].contiguous()
+
+ loss = self.loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1))
+
+ return loss
+
+
+class MixFormerSequentialPreTrainedModel(PreTrainedModel):
+ """MixFormer (sequential for DeepSpeed) pre-trained model."""
+
+ config_class = MixFormerSequentialConfig
+ base_model_prefix = "transformer"
+ supports_gradient_checkpointing = True
+
+ def __init__(self, *inputs, **kwargs) -> None:
+ super().__init__(*inputs, **kwargs)
+
+ def _init_weights(self, module: nn.Module) -> None:
+ if isinstance(module, (nn.Linear,)):
+ module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+ if module.bias is not None:
+ module.bias.data.zero_()
+ elif isinstance(module, nn.Embedding):
+ module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+ if module.padding_idx is not None:
+ module.weight.data[module.padding_idx].zero_()
+ elif isinstance(module, nn.LayerNorm):
+ module.bias.data.zero_()
+ module.weight.data.fill_(1.0)
+
+ def prepare_inputs_for_generation(
+ self,
+ input_ids: torch.LongTensor,
+ past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,
+ attention_mask: Optional[torch.BoolTensor] = None,
+ **kwargs,
+ ) -> Dict[str, Any]:
+ if attention_mask is not None and torch.any(~attention_mask.bool()):
+ total_seq_len = torch.sum(attention_mask, dim=1)
+ max_seq_len = torch.max(total_seq_len)
+
+ total_seq_len = torch.cat((torch.tensor([0], device=attention_mask.device), total_seq_len)).unsqueeze(1)
+ cumulative_seq_len = torch.cumsum(total_seq_len, dim=0).squeeze(1).to(torch.int32)
+ attention_mask = (attention_mask.bool(), cumulative_seq_len, max_seq_len.item())
+ else:
+ attention_mask = None
+
+ if past_key_values is None or not (isinstance(past_key_values, InferenceParams)):
+ past_key_values = InferenceParams(
+ max_batch_size=input_ids.shape[0],
+ max_sequence_len=self.config.n_positions,
+ sequence_len_offset=0,
+ batch_size_offset=0,
+ fused_ft_kernel=False,
+ key_value_memory_dict={},
+ )
+ else:
+ # Assume that `past_key_values` has cached all tokens up to the last token in `input_ids`
+ past_key_values.sequence_len_offset = len(input_ids[0]) - 1
+ input_ids = input_ids[:, -1].unsqueeze(-1)
+
+ return {
+ "input_ids": input_ids,
+ "past_key_values": past_key_values,
+ "attention_mask": attention_mask,
+ }
+
+ # Copied from transformers.models.gpt2.modeling_gpt2.GPT2PreTrainedModel._set_gradient_checkpointing with GPT2->GPTBigCode
+ def _set_gradient_checkpointing(self, module, value=False):
+ if isinstance(module, PreTrainedModel):
+ module.gradient_checkpointing = value
+
+
+class MixFormerSequentialForCausalLM(MixFormerSequentialPreTrainedModel):
+ """MixFormer (sequential for DeepSpeed) for Causal Language Modeling."""
+
+ _keys_to_ignore_on_load_missing = [""]
+ _keys_to_ignore_on_load_unexpected = [r"layers\.\d+\.mlp.(fc_in|fc_out)\.(weight|bias)"]
+ _no_split_modules = ["ParallelBlock"]
+
+ def __init__(self, config: MixFormerSequentialConfig) -> None:
+ super().__init__(config)
+
+ modules = [Embedding(config)]
+ modules += [ParallelBlock(config, block_idx=i) for i in range(config.n_layer)]
+ modules.append(CausalLMHead(config))
+
+ self.layers = nn.Sequential(*modules)
+ self.loss = CausalLMLoss()
+
+ self.post_init()
+
+ def get_input_embeddings(self) -> nn.Embedding:
+ return self.layers[0].wte
+
+ def set_input_embeddings(self, new_embeddings: nn.Embedding) -> None:
+ self.layers[0].wte = new_embeddings
+
+ def get_output_embeddings(self) -> nn.Linear:
+ return self.layers[-1].linear
+
+ def set_output_embeddings(self, new_embeddings: nn.Linear) -> None:
+ self.layers[-1].linear = new_embeddings
+
+ def forward(
+ self,
+ input_ids: torch.LongTensor,
+ past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,
+ attention_mask: Optional[torch.BoolTensor] = None,
+ labels: Optional[torch.LongTensor] = None,
+ **kwargs,
+ ) -> CausalLMOutputWithPast:
+ if attention_mask is not None and self.training:
+ print("`attention_mask` is not supported during training. Using it might lead to unexpected results.")
+
+ if past_key_values is None and attention_mask is None:
+ lm_logits = self.layers(input_ids)
+ else:
+ hidden_layer = self.layers[0](input_ids)
+ for module in self.layers[1:-1]:
+ hidden_layer = module(hidden_layer, past_key_values=past_key_values, attention_mask=attention_mask)
+ lm_logits = self.layers[-1](hidden_layer)
+
+ loss = None
+ if labels is not None:
+ loss = self.loss(lm_logits, labels)
+
+ return CausalLMOutputWithPast(loss=loss, logits=lm_logits, past_key_values=past_key_values)
diff --git a/mftcoder_accelerate/src/model/qwen/cache_autogptq_cuda_256.cpp b/mftcoder_accelerate/src/model/qwen/cache_autogptq_cuda_256.cpp
new file mode 100644
index 0000000..8458a9b
--- /dev/null
+++ b/mftcoder_accelerate/src/model/qwen/cache_autogptq_cuda_256.cpp
@@ -0,0 +1,198 @@
+#include
+#include
+#include
+
+// adapted from https://github.com/PanQiWei/AutoGPTQ/blob/main/autogptq_extension/cuda_256/autogptq_cuda_256.cpp
+void vecquant8matmul_cuda(
+ torch::Tensor vec, torch::Tensor mat, torch::Tensor mul,
+ torch::Tensor scales, torch::Tensor zeros,
+ torch::Tensor g_idx
+);
+
+void vecquant8matmul(
+ torch::Tensor vec, torch::Tensor mat, torch::Tensor mul,
+ torch::Tensor scales, torch::Tensor zeros,
+ torch::Tensor g_idx
+) {
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(vec));
+ vecquant8matmul_cuda(vec, mat, mul, scales, zeros, g_idx);
+}
+
+void vecquant8matmul_batched_cuda(
+ torch::Tensor vec, torch::Tensor mat, torch::Tensor mul,
+ torch::Tensor scales, torch::Tensor zeros
+);
+
+void vecquant8matmul_batched(
+ torch::Tensor vec, torch::Tensor mat, torch::Tensor mul,
+ torch::Tensor scales, torch::Tensor zeros
+) {
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(vec));
+ vecquant8matmul_batched_cuda(vec, mat, mul, scales, zeros);
+}
+
+void vecquant8matmul_batched_column_compression_cuda(
+ torch::Tensor vec, torch::Tensor mat, torch::Tensor mul,
+ torch::Tensor scales, torch::Tensor zeros
+);
+
+void vecquant8matmul_batched_column_compression(
+ torch::Tensor vec, torch::Tensor mat, torch::Tensor mul,
+ torch::Tensor scales, torch::Tensor zeros
+) {
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(vec));
+ vecquant8matmul_batched_column_compression_cuda(vec, mat, mul, scales, zeros);
+}
+
+void vecquant4matmul_batched_cuda(
+ torch::Tensor vec, torch::Tensor mat, torch::Tensor mul,
+ torch::Tensor scales, torch::Tensor zeros
+);
+
+void vecquant4matmul_batched(
+ torch::Tensor vec, torch::Tensor mat, torch::Tensor mul,
+ torch::Tensor scales, torch::Tensor zeros
+) {
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(vec));
+ vecquant4matmul_batched_cuda(vec, mat, mul, scales, zeros);
+}
+
+void vecquant4matmul_batched_column_compression_cuda(
+ torch::Tensor vec, torch::Tensor mat, torch::Tensor mul,
+ torch::Tensor scales, torch::Tensor zeros
+);
+
+void vecquant4matmul_batched_column_compression(
+ torch::Tensor vec, torch::Tensor mat, torch::Tensor mul,
+ torch::Tensor scales, torch::Tensor zeros
+) {
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(vec));
+ vecquant4matmul_batched_column_compression_cuda(vec, mat, mul, scales, zeros);
+}
+
+void vecquant8matmul_batched_old_cuda(
+ torch::Tensor vec, torch::Tensor mat, torch::Tensor mul,
+ torch::Tensor scales, torch::Tensor zeros
+);
+
+void vecquant8matmul_batched_old(
+ torch::Tensor vec, torch::Tensor mat, torch::Tensor mul,
+ torch::Tensor scales, torch::Tensor zeros
+) {
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(vec));
+ vecquant8matmul_batched_old_cuda(vec, mat, mul, scales, zeros);
+}
+
+
+void vecquant4matmul_batched_old_cuda(
+ torch::Tensor vec, torch::Tensor mat, torch::Tensor mul,
+ torch::Tensor scales, torch::Tensor zeros
+);
+
+void vecquant4matmul_batched_old(
+ torch::Tensor vec, torch::Tensor mat, torch::Tensor mul,
+ torch::Tensor scales, torch::Tensor zeros
+) {
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(vec));
+ vecquant4matmul_batched_old_cuda(vec, mat, mul, scales, zeros);
+}
+
+void vecquant8matmul_batched_column_compression_old_cuda(
+ torch::Tensor vec, torch::Tensor mat, torch::Tensor mul,
+ torch::Tensor scales, torch::Tensor zeros
+);
+
+void vecquant8matmul_batched_column_compression_old(
+ torch::Tensor vec, torch::Tensor mat, torch::Tensor mul,
+ torch::Tensor scales, torch::Tensor zeros
+) {
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(vec));
+ vecquant8matmul_batched_column_compression_old_cuda(vec, mat, mul, scales, zeros);
+}
+
+void vecquant4matmul_batched_column_compression_old_cuda(
+ torch::Tensor vec, torch::Tensor mat, torch::Tensor mul,
+ torch::Tensor scales, torch::Tensor zeros
+);
+
+void vecquant4matmul_batched_column_compression_old(
+ torch::Tensor vec, torch::Tensor mat, torch::Tensor mul,
+ torch::Tensor scales, torch::Tensor zeros
+) {
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(vec));
+ vecquant4matmul_batched_column_compression_old_cuda(vec, mat, mul, scales, zeros);
+}
+
+
+
+void vecquant8matmul_batched_faster_cuda(
+ torch::Tensor vec, torch::Tensor mat, torch::Tensor mul,
+ torch::Tensor scales, torch::Tensor zeros
+);
+
+void vecquant8matmul_batched_faster(
+ torch::Tensor vec, torch::Tensor mat, torch::Tensor mul,
+ torch::Tensor scales, torch::Tensor zeros
+) {
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(vec));
+ vecquant8matmul_batched_faster_cuda(vec, mat, mul, scales, zeros);
+}
+
+
+void vecquant8matmul_batched_faster_old_cuda(
+ torch::Tensor vec, torch::Tensor mat, torch::Tensor mul,
+ torch::Tensor scales, torch::Tensor zeros
+);
+
+void vecquant8matmul_batched_faster_old(
+ torch::Tensor vec, torch::Tensor mat, torch::Tensor mul,
+ torch::Tensor scales, torch::Tensor zeros
+) {
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(vec));
+ vecquant8matmul_batched_faster_old_cuda(vec, mat, mul, scales, zeros);
+}
+
+void vecquant8matmul_batched_column_compression_faster_cuda(
+ torch::Tensor vec, torch::Tensor mat, torch::Tensor mul,
+ torch::Tensor scales, torch::Tensor zeros
+);
+
+void vecquant8matmul_batched_column_compression_faster(
+ torch::Tensor vec, torch::Tensor mat, torch::Tensor mul,
+ torch::Tensor scales, torch::Tensor zeros
+) {
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(vec));
+ vecquant8matmul_batched_column_compression_faster_cuda(vec, mat, mul, scales, zeros);
+}
+
+
+void vecquant8matmul_batched_column_compression_faster_old_cuda(
+ torch::Tensor vec, torch::Tensor mat, torch::Tensor mul,
+ torch::Tensor scales, torch::Tensor zeros
+);
+
+void vecquant8matmul_batched_column_compression_faster_old(
+ torch::Tensor vec, torch::Tensor mat, torch::Tensor mul,
+ torch::Tensor scales, torch::Tensor zeros
+) {
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(vec));
+ vecquant8matmul_batched_column_compression_faster_old_cuda(vec, mat, mul, scales, zeros);
+}
+
+
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+ m.def("vecquant8matmul", &vecquant8matmul, "Vector 8-bit Quantized Matrix Multiplication (CUDA) (desc_act)");
+ m.def("vecquant8matmul_batched", &vecquant8matmul_batched, "Vector 8-bit Batched Quantized Matrix Multiplication (CUDA) (desc_act)");
+ m.def("vecquant8matmul_batched_old", &vecquant8matmul_batched_old, "Vector 8-bit old Batched Quantized Matrix Multiplication (CUDA) (desc_act)");
+ m.def("vecquant8matmul_batched_faster", &vecquant8matmul_batched_faster, "Vector 8-bit old Batched Quantized Matrix Multiplication (CUDA) (desc_act)");
+ m.def("vecquant8matmul_batched_faster_old", &vecquant8matmul_batched_faster_old, "Vector 8-bit old Batched Quantized Matrix Multiplication (CUDA) (desc_act)");
+ m.def("vecquant4matmul_batched_old", &vecquant4matmul_batched_old, "Vector 4-bit old Batched Quantized Matrix Multiplication (CUDA) (desc_act)");
+ m.def("vecquant8matmul_batched_column_compression", &vecquant8matmul_batched_column_compression, "Vector 8-bit Batched Quantized Matrix Multiplication (CUDA) with weight's column compressed (desc_act)");
+ m.def("vecquant8matmul_batched_column_compression_old", &vecquant8matmul_batched_column_compression_old, "Vector old 8-bit Batched Quantized Matrix Multiplication (CUDA) with weight's column compressed (desc_act)");
+ m.def("vecquant8matmul_batched_column_compression_faster", &vecquant8matmul_batched_column_compression_faster, "Vector old 8-bit Batched Quantized Matrix Multiplication (CUDA) with weight's column compressed (desc_act)");
+ m.def("vecquant8matmul_batched_column_compression_faster_old", &vecquant8matmul_batched_column_compression_faster_old, "Vector old 8-bit Batched Quantized Matrix Multiplication (CUDA) with weight's column compressed (desc_act)");
+ m.def("vecquant4matmul_batched_column_compression_old", &vecquant4matmul_batched_column_compression_old, "Vector old 4-bit Batched Quantized Matrix Multiplication (CUDA) with weight's column compressed (desc_act)");
+ m.def("vecquant4matmul_batched", &vecquant4matmul_batched, "Vector 4-bit Batched Quantized Matrix Multiplication (CUDA) (desc_act)");
+ m.def("vecquant4matmul_batched_column_compression", &vecquant4matmul_batched_column_compression, "Vector 4-bit Batched Quantized Matrix Multiplication (CUDA) with weight's column compressed (desc_act)");
+}
diff --git a/mftcoder_accelerate/src/model/qwen/cache_autogptq_cuda_kernel_256.cu b/mftcoder_accelerate/src/model/qwen/cache_autogptq_cuda_kernel_256.cu
new file mode 100644
index 0000000..b7932cd
--- /dev/null
+++ b/mftcoder_accelerate/src/model/qwen/cache_autogptq_cuda_kernel_256.cu
@@ -0,0 +1,1708 @@
+#define _CRT_SECURE_NO_WARNINGS
+#include
+#include
+#include
+#include
+#include
+#include
+
+#if (defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 700) || defined(USE_ROCM)
+// adapted from https://github.com/PanQiWei/AutoGPTQ/blob/main/autogptq_extension/cuda_256/autogptq_cuda_kernel_256.cu
+__device__ __forceinline__ void atomicAdd(c10::Half* address, c10::Half val) {
+ unsigned int *address_as_ui = reinterpret_cast(reinterpret_cast(address) - (reinterpret_cast(address) & 2));
+ unsigned int old = *address_as_ui;
+ unsigned int assumed;
+
+ do {
+ assumed = old;
+ unsigned short hsum = reinterpret_cast(address) & 2 ? (old >> 16) : (old & 0xffff);
+ hsum += val;
+ old = reinterpret_cast(address) & 2
+ ? (old & 0xffff) | (hsum << 16)
+ : (old & 0xffff0000) | hsum;
+ old = atomicCAS(address_as_ui, assumed, old);
+
+ // Note: uses integer comparison to avoid hang in case of NaN (since NaN != NaN)
+ } while (assumed != old);
+}
+__device__ __forceinline__ void atomicAdd(__half* address, c10::Half val) {
+ unsigned int * address_as_ui = (unsigned int *) ((char *)address - ((size_t)address & 2));
+ unsigned int old = *address_as_ui;
+ unsigned int assumed;
+
+ do {
+ assumed = old;
+ __half_raw hsum;
+ hsum.x = (size_t)address & 2 ? (old >> 16) : (old & 0xffff);
+ half tmpres = __hadd(hsum, val);
+ hsum = __half_raw(tmpres);
+ old = (size_t)address & 2 ? (old & 0xffff) | (hsum.x << 16) : (old & 0xffff0000) | hsum.x;
+ old = atomicCAS(address_as_ui, assumed, old);
+ } while (assumed != old);
+}
+#endif
+
+template
+__global__ void VecQuant8MatMulKernel(
+ const scalar_t* __restrict__ vec,
+ const int* __restrict__ mat,
+ scalar_t* __restrict__ mul,
+ const scalar_t* __restrict__ scales,
+ const int* __restrict__ zeros,
+ const int* __restrict__ g_idx,
+ int batch,
+ int vec_height,
+ int height,
+ int width,
+ int zero_width
+);
+
+template
+__global__ void VecQuant8BatchMatMulColumnCompressionKernel(
+ const scalar_t* __restrict__ vec,
+ const int* __restrict__ mat,
+ scalar_t* __restrict__ mul,
+ const scalar_t* __restrict__ scales,
+ const int* __restrict__ zeros,
+ int batch,
+ int heads,
+ int vec_row,
+ int height,
+ int width
+);
+
+template
+__global__ void VecQuant4BatchMatMulColumnCompressionKernel(
+ const scalar_t* __restrict__ vec,
+ const int* __restrict__ mat,
+ scalar_t* __restrict__ mul,
+ const scalar_t* __restrict__ scales,
+ const int* __restrict__ zeros,
+ int batch,
+ int heads,
+ int vec_row,
+ int height,
+ int width
+);
+
+template
+__global__ void VecQuant8BatchMatMulKernel(
+ const scalar_t* __restrict__ vec,
+ const int* __restrict__ mat,
+ scalar_t* __restrict__ mul,
+ const scalar_t* __restrict__ scales,
+ const int* __restrict__ zeros,
+ int batch,
+ int heads,
+ int vec_row,
+ int vec_height,
+ int height,
+ int width,
+ int zero_width
+);
+
+template
+__global__ void VecQuant4BatchMatMulKernel(
+ const scalar_t* __restrict__ vec,
+ const int* __restrict__ mat,
+ scalar_t* __restrict__ mul,
+ const scalar_t* __restrict__ scales,
+ const int* __restrict__ zeros,
+ int batch,
+ int heads,
+ int vec_row,
+ int vec_height,
+ int height,
+ int width,
+ int zero_width
+);
+
+
+
+template
+__global__ void VecQuant8BatchMatMulKernel_old(
+ const scalar_t* __restrict__ vec,
+ const uint8_t* __restrict__ mat,
+ scalar_t* __restrict__ mul,
+ const scalar_t* __restrict__ scales,
+ const scalar_t* __restrict__ zeros,
+ int batch,
+ int heads,
+ int vec_row,
+ int vec_height,
+ int height,
+ int width,
+ int zero_width
+);
+
+__global__ void VecQuant8BatchMatMulKernel_faster(
+ const half* __restrict__ vec,
+ const uint8_t* __restrict__ mat,
+ half* __restrict__ mul,
+ const half* __restrict__ scales,
+ const half* __restrict__ zeros,
+ int batch,
+ int heads,
+ int vec_row,
+ int vec_height,
+ int height,
+ int width,
+ int zero_width
+);
+
+
+
+__global__ void VecQuant8BatchMatMulKernel_faster_old(
+ const half* __restrict__ vec,
+ const uint8_t* __restrict__ mat,
+ half* __restrict__ mul,
+ const half* __restrict__ scales,
+ const half* __restrict__ zeros,
+ int batch,
+ int heads,
+ int vec_row,
+ int vec_height,
+ int height,
+ int width
+);
+
+
+template
+__global__ void VecQuant4BatchMatMulKernel_old(
+ const scalar_t* __restrict__ vec,
+ const uint8_t* __restrict__ mat,
+ scalar_t* __restrict__ mul,
+ const scalar_t* __restrict__ scales,
+ const scalar_t* __restrict__ zeros,
+ int batch,
+ int heads,
+ int vec_row,
+ int vec_height,
+ int height,
+ int width,
+ int zero_width
+);
+
+
+template
+__global__ void VecQuant8BatchMatMulColumnCompressionKernel_old(
+ const scalar_t* __restrict__ vec,
+ const uint8_t* __restrict__ mat,
+ scalar_t* __restrict__ mul,
+ const scalar_t* __restrict__ scales,
+ const scalar_t* __restrict__ zeros,
+ int batch,
+ int heads,
+ int vec_row,
+ int height,
+ int width
+);
+
+__global__ void VecQuant8BatchMatMulColumnCompressionKernel_faster(
+ const half* __restrict__ vec,
+ const uint8_t* __restrict__ mat,
+ half* __restrict__ mul,
+ const half* __restrict__ scales,
+ const half* __restrict__ zeros,
+ int batch,
+ int heads,
+ int vec_row,
+ int height,
+ int width
+);
+
+__global__ void VecQuant8BatchMatMulColumnCompressionKernel_faster_old(
+ const half* __restrict__ vec,
+ const uint8_t* __restrict__ mat,
+ half* __restrict__ mul,
+ const half* __restrict__ scales,
+ const half* __restrict__ zeros,
+ int batch,
+ int heads,
+ int vec_row,
+ int height,
+ int width
+);
+
+
+template
+__global__ void VecQuant4BatchMatMulColumnCompressionKernel_old(
+ const scalar_t* __restrict__ vec,
+ const uint8_t* __restrict__ mat,
+ scalar_t* __restrict__ mul,
+ const scalar_t* __restrict__ scales,
+ const scalar_t* __restrict__ zeros,
+ int batch,
+ int heads,
+ int vec_row,
+ int height,
+ int width
+);
+
+
+__global__ void VecQuant8BatchMatMulKernel_faster(
+ const half* __restrict__ vec,
+ const uint8_t* __restrict__ mat,
+ half* __restrict__ mul,
+ const half* __restrict__ scales,
+ const half* __restrict__ zeros,
+ int batch,
+ int heads,
+ int vec_row,
+ int vec_height,
+ int height,
+ int width
+);
+
+
+__global__ void VecQuant8BatchMatMulColumnCompressionKernel_faster(
+ const half* __restrict__ vec,
+ const uint8_t* __restrict__ mat,
+ half* __restrict__ mul,
+ const half* __restrict__ scales,
+ const half* __restrict__ zeros,
+ int batch,
+ int heads,
+ int vec_row,
+ int height,
+ int width
+);
+
+const int BLOCKWIDTH = 128;
+const int BLOCKHEIGHT8 = 32;
+const int BLOCKHEIGHT4 = 16;
+const int BLOCKHEIGHT_OLD4 = 128;
+//const int BLOCKHEIGHT_OLD8 = 128;
+
+__device__ inline unsigned int as_unsigned(int i) {
+ return *reinterpret_cast(&i);
+}
+
+__device__ inline int as_int(int i) {
+ return *reinterpret_cast(&i);
+}
+
+void vecquant8matmul_batched_column_compression_cuda(
+ torch::Tensor vec,
+ torch::Tensor mat,
+ torch::Tensor mul,
+ torch::Tensor scales,
+ torch::Tensor zeros
+) {
+ int batch = vec.size(0);
+ int heads = vec.size(1);
+ int vec_row = vec.size(2);
+ int height = vec.size(3);
+ int width = mat.size(3) * 4;
+
+ dim3 blocks(
+ (height + BLOCKWIDTH - 1) / BLOCKWIDTH,
+ (width + BLOCKWIDTH - 1) / BLOCKWIDTH
+ );
+ dim3 threads(BLOCKWIDTH);
+
+ AT_DISPATCH_FLOATING_TYPES(
+ vec.type(), "vecquant8matmul_batched_cuda", ([&] {
+ VecQuant8BatchMatMulColumnCompressionKernel<<>>(
+ vec.data(), mat.data(), mul.data(),
+ scales.data(), zeros.data(),
+ batch, heads, vec_row, height, width
+ );
+ })
+ );
+
+}
+
+template
+__global__ void VecQuant8BatchMatMulColumnCompressionKernel(
+ const scalar_t* __restrict__ vec,
+ const int* __restrict__ mat,
+ scalar_t* __restrict__ mul,
+ const scalar_t* __restrict__ scales,
+ const int* __restrict__ zeros,
+ int batch,
+ int heads,
+ int vec_row,
+ int height,
+ int width
+) {
+ int weight_total = batch * heads * height * width / 4;
+ int input_total = batch * heads * vec_row * height;
+ int out_total = batch * heads * vec_row * width;
+ int tid = threadIdx.x;
+ // h is index of height with step being BLOCKWIDTH
+ int h = BLOCKWIDTH * blockIdx.x;
+ // w is index of width with step being 1
+ int w = BLOCKWIDTH * blockIdx.y + tid;
+ if (w >= width && tid >= height) {
+ return;
+ }
+
+ __shared__ scalar_t blockvec[BLOCKWIDTH];
+ int k;
+ scalar_t w_tmp;
+
+ float weight[BLOCKWIDTH];
+
+ for (int b = 0; b < batch; ++b){
+ for (int head = 0; head < heads; ++head){
+ int batch_shift = b * heads + head;
+ for (k = 0; k < BLOCKWIDTH && h + k < height; ++k){
+ int i_w = (w / 4);
+ int w_bit = (w % 4) * 8;
+
+ int w_index = (batch_shift * height + h + k) * width / 4 + i_w;
+ if (w_index >= weight_total || w >= width) {
+ weight[k] = 0;
+ } else {
+ scalar_t scale = scales[batch_shift * height + h + k];
+ scalar_t zero = zeros[batch_shift * height + h + k];
+ w_tmp = ((as_unsigned(mat[w_index]) >> w_bit) & 0xFF);
+ weight[k] = scale * (w_tmp - zero);
+ }
+ }
+
+ scalar_t res;
+ for (int vr = 0; vr < vec_row; ++vr){
+ res = 0;
+ int vec_index = (batch_shift * vec_row + vr) * height + blockIdx.x * BLOCKWIDTH + tid;
+ if (vec_index < input_total) {
+ blockvec[tid] = vec[vec_index];
+ } else {
+ blockvec[tid] = 0;
+ }
+
+ __syncthreads();
+ for (k = 0; k < BLOCKWIDTH && h + k < height; ++k){
+ // res is the dot product of BLOCKWIDTH elements (part of width)
+ res += weight[k] * blockvec[k];
+ }
+ // add res to the final result, final matrix shape: (batch, vec_row, width)
+ int out_index = (batch_shift * vec_row + vr) * width + w;
+ if (out_index < out_total) {
+ atomicAdd(&mul[out_index], res);
+ }
+ __syncthreads();
+ }
+ }
+ }
+}
+
+void vecquant8matmul_batched_cuda(
+ torch::Tensor vec,
+ torch::Tensor mat,
+ torch::Tensor mul,
+ torch::Tensor scales,
+ torch::Tensor zeros
+) {
+ int batch = vec.size(0);
+ int heads = vec.size(1);
+ int vec_row = vec.size(2);
+ int vec_height = vec.size(3);
+ int height = mat.size(2);
+ int width = mat.size(3);
+ int zero_width = zeros.size(2);
+
+ dim3 blocks(
+ (height + BLOCKHEIGHT8 - 1) / BLOCKHEIGHT8,
+ (width + BLOCKWIDTH - 1) / BLOCKWIDTH
+ );
+ dim3 threads(BLOCKWIDTH);
+
+ AT_DISPATCH_FLOATING_TYPES(
+ vec.type(), "vecquant8matmul_batched_cuda", ([&] {
+ VecQuant8BatchMatMulKernel<<>>(
+ vec.data(), mat.data(), mul.data(),
+ scales.data(), zeros.data(),
+ batch, heads, vec_row, vec_height, height, width, zero_width
+ );
+ })
+ );
+
+}
+
+template
+__global__ void VecQuant8BatchMatMulKernel(
+ const scalar_t* __restrict__ vec,
+ const int* __restrict__ mat,
+ scalar_t* __restrict__ mul,
+ const scalar_t* __restrict__ scales,
+ const int* __restrict__ zeros,
+ int batch,
+ int heads,
+ int vec_row,
+ int vec_height,
+ int height,
+ int width,
+ int zero_width
+) {
+ int weight_total = batch * heads * height * width;
+ int input_total = batch * heads * vec_row * vec_height;
+ int out_total = batch * heads * vec_row * width;
+ int tid = threadIdx.x;
+ // h is index of height with step being BLOCKHEIGHT8
+ int h = BLOCKHEIGHT8 * blockIdx.x;
+ // w is index of width with step being 1
+ int w = BLOCKWIDTH * blockIdx.y + tid;
+ if (w >= width && tid >= vec_height) {
+ return;
+ }
+
+ __shared__ scalar_t blockvec[BLOCKWIDTH];
+ // i is index of mat of block first row
+ int i = width * h + w;
+ // if (i >= width * height) {
+ // return;
+ // }
+ int k;
+ scalar_t w_tmp;
+
+ int z_w = w / 4;
+ int z_mod = (w % 4) * 8;
+
+ float weight[BLOCKWIDTH];
+
+ for (int b = 0; b < batch; ++b){
+ for (int head = 0; head < heads; ++head){
+ int batch_shift = b * heads + head;
+ for (k = 0; k < BLOCKWIDTH && h * 4 + k < vec_height; ++k){
+ int k_w = (k / 4);
+ int k_bit = (k % 4) * 8;
+
+ int w_index = batch_shift * height * width + i + (k_w * width);
+ if (w_index >= weight_total || w >= width) {
+ weight[k] = 0;
+ } else {
+ scalar_t scale = scales[batch_shift * width + w];
+ scalar_t zero;
+ if (zero_width == width) {
+ zero = zeros[batch_shift * width + w];
+ } else {
+ zero = scalar_t(((as_unsigned(zeros[batch_shift * zero_width + z_w]) >> z_mod) & 0xFF) + 1);
+ }
+ w_tmp = ((as_unsigned(mat[w_index]) >> k_bit) & 0xFF);
+ weight[k] = scale * (w_tmp - zero);
+ }
+ }
+
+ scalar_t res;
+ for (int vr = 0; vr < vec_row; ++vr){
+ res = 0;
+ int vec_index = (batch_shift * vec_row + vr) * vec_height + blockIdx.x * BLOCKWIDTH + tid;
+ if (vec_index < input_total) {
+ blockvec[tid] = vec[vec_index];
+ } else {
+ blockvec[tid] = 0;
+ }
+
+ __syncthreads();
+ for (k = 0; k < BLOCKWIDTH && h * 4 + k < vec_height; ++k){
+ // res is the dot product of BLOCKWIDTH elements (part of width)
+ res += weight[k] * blockvec[k];
+ }
+ // add res to the final result, final matrix shape: (batch, vec_row, width)
+ int out_index = (batch_shift * vec_row + vr) * width + w;
+ if (out_index < out_total) {
+ atomicAdd(&mul[out_index], res);
+ }
+ __syncthreads();
+ }
+ }
+ }
+}
+
+
+void vecquant8matmul_cuda(
+ torch::Tensor vec,
+ torch::Tensor mat,
+ torch::Tensor mul,
+ torch::Tensor scales,
+ torch::Tensor zeros,
+ torch::Tensor g_idx
+) {
+ int batch = vec.size(0);
+ int vec_height = vec.size(1);
+ int height = mat.size(0);
+ int width = mat.size(1);
+ int zero_width = zeros.size(1);
+
+ dim3 blocks(
+ (height + BLOCKHEIGHT8 - 1) / BLOCKHEIGHT8,
+ (width + BLOCKWIDTH - 1) / BLOCKWIDTH
+ );
+ dim3 threads(BLOCKWIDTH);
+
+ AT_DISPATCH_FLOATING_TYPES(
+ vec.type(), "vecquant8matmul_cuda", ([&] {
+ VecQuant8MatMulKernel<<>>(
+ vec.data(), mat.data(), mul.data(),
+ scales.data(), zeros.data(), g_idx.data(),
+ batch, vec_height, height, width, zero_width
+ );
+ })
+ );
+}
+
+template
+__global__ void VecQuant8MatMulKernel(
+ const scalar_t* __restrict__ vec,
+ const int* __restrict__ mat,
+ scalar_t* __restrict__ mul,
+ const scalar_t* __restrict__ scales,
+ const int* __restrict__ zeros,
+ const int* __restrict__ g_idx,
+ int batch,
+ int vec_height,
+ int height,
+ int width,
+ int zero_width
+) {
+ int h = BLOCKHEIGHT8 * blockIdx.x;
+ int w = BLOCKWIDTH * blockIdx.y + threadIdx.x;
+
+ __shared__ scalar_t blockvec[BLOCKWIDTH];
+ int i = width * h + w;
+ int g_h = h * 4;
+ int k;
+ unsigned int g;
+ scalar_t w_tmp;
+
+ int z_w = w / 4;
+ int z_mod = (w % 4) * 8;
+
+ float weight[BLOCKWIDTH];
+
+ for (k = 0; k < BLOCKWIDTH; ++k){
+ int k_w = (k / 4);
+ int k_bit = (k % 4) * 8;
+
+ g = as_int(g_idx[g_h + k]);
+ scalar_t scale = scales[g * width + w];
+ scalar_t zero = scalar_t(((as_unsigned(zeros[g * zero_width + z_w]) >> z_mod) & 0xFF) + 1);
+
+ w_tmp = ((as_unsigned(mat[i + (k_w * width)]) >> k_bit) & 0xFF);
+
+ weight[k] = scale * (w_tmp - zero);
+ }
+
+
+ scalar_t res;
+ for (int b = 0; b < batch; ++b){
+ res = 0;
+ blockvec[threadIdx.x] = vec[b * vec_height + blockIdx.x * BLOCKWIDTH + threadIdx.x];
+ __syncthreads();
+ for (k = 0; k < BLOCKWIDTH; ++k){
+ res += weight[k] * blockvec[k];
+ }
+ atomicAdd(&mul[b * width + w], res);
+ __syncthreads();
+ }
+}
+
+
+
+void vecquant4matmul_batched_cuda(
+ torch::Tensor vec,
+ torch::Tensor mat,
+ torch::Tensor mul,
+ torch::Tensor scales,
+ torch::Tensor zeros
+) {
+ int batch = vec.size(0);
+ int heads = vec.size(1);
+ int vec_row = vec.size(2);
+ int vec_height = vec.size(3);
+ int height = mat.size(2);
+ int width = mat.size(3);
+ int zero_width = zeros.size(2);
+
+ dim3 blocks(
+ (height + BLOCKHEIGHT4 - 1) / BLOCKHEIGHT4,
+ (width + BLOCKWIDTH - 1) / BLOCKWIDTH
+ );
+ dim3 threads(BLOCKWIDTH);
+
+ AT_DISPATCH_FLOATING_TYPES(
+ vec.type(), "vecquant4matmul_batched_cuda", ([&] {
+ VecQuant4BatchMatMulKernel<<>>(
+ vec.data(), mat.data(), mul.data(),
+ scales.data(), zeros.data(),
+ batch, heads, vec_row, vec_height, height, width, zero_width
+ );
+ })
+ );
+
+}
+
+template
+__global__ void VecQuant4BatchMatMulKernel(
+ const scalar_t* __restrict__ vec,
+ const int* __restrict__ mat,
+ scalar_t* __restrict__ mul,
+ const scalar_t* __restrict__ scales,
+ const int* __restrict__ zeros,
+ int batch,
+ int heads,
+ int vec_row,
+ int vec_height,
+ int height,
+ int width,
+ int zero_width
+) {
+ int weight_total = batch * heads * height * width;
+ int input_total = batch * heads * vec_row * vec_height;
+ int out_total = batch * heads * vec_row * width;
+ int tid = threadIdx.x;
+ // h is index of height with step being BLOCKHEIGHT4
+ int h = BLOCKHEIGHT4 * blockIdx.x;
+ // w is index of width with step being 1
+ int w = BLOCKWIDTH * blockIdx.y + tid;
+ if (w >= width && tid >= vec_height) {
+ return;
+ }
+
+ __shared__ scalar_t blockvec[BLOCKWIDTH];
+ // i is index of mat of block first row
+ int i = width * h + w;
+ int k;
+ scalar_t w_tmp;
+
+ int z_w = w / 8;
+ int z_mod = (w % 8) * 4;
+
+ float weight[BLOCKWIDTH];
+
+ for (int b = 0; b < batch; ++b){
+ for (int head = 0; head < heads; ++head){
+ int batch_shift = b * heads + head;
+ for (k = 0; k < BLOCKWIDTH && h * 8 + k < vec_height; ++k){
+ int k_w = (k / 8);
+ int k_bit = (k % 8) * 4;
+
+ int w_index = batch_shift * height * width + i + (k_w * width);
+ if (w_index >= weight_total || w >= width) {
+ weight[k] = 0;
+ } else {
+ scalar_t scale = scales[batch_shift * width + w];
+ scalar_t zero;
+ if (zero_width == width) {
+ zero = zeros[batch_shift * width + w];
+ } else {
+ zero = scalar_t(((as_unsigned(zeros[batch_shift * zero_width + z_w]) >> z_mod) & 0xF));
+ }
+ w_tmp = ((as_unsigned(mat[w_index]) >> k_bit) & 0xF);
+ weight[k] = scale * (w_tmp - zero);
+ }
+ }
+
+ scalar_t res;
+ for (int vr = 0; vr < vec_row; ++vr){
+ res = 0;
+ int vec_index = (batch_shift * vec_row + vr) * vec_height + blockIdx.x * BLOCKWIDTH + tid;
+ if (vec_index < input_total) {
+ blockvec[tid] = vec[vec_index];
+ } else {
+ blockvec[tid] = 0;
+ }
+
+ __syncthreads();
+ for (k = 0; k < BLOCKWIDTH && h * 8 + k < vec_height; ++k){
+ // res is the dot product of BLOCKWIDTH elements (part of width)
+ res += weight[k] * blockvec[k];
+ }
+ // add res to the final result, final matrix shape: (batch, vec_row, width)
+ int out_index = (batch_shift * vec_row + vr) * width + w;
+ if (out_index < out_total) {
+ atomicAdd(&mul[out_index], res);
+ }
+ __syncthreads();
+ }
+ }
+ }
+}
+
+
+
+void vecquant4matmul_batched_column_compression_cuda(
+ torch::Tensor vec,
+ torch::Tensor mat,
+ torch::Tensor mul,
+ torch::Tensor scales,
+ torch::Tensor zeros
+) {
+ int batch = vec.size(0);
+ int heads = vec.size(1);
+ int vec_row = vec.size(2);
+ int height = vec.size(3);
+ int width = mat.size(3) * 8;
+
+ dim3 blocks(
+ (height + BLOCKWIDTH - 1) / BLOCKWIDTH,
+ (width + BLOCKWIDTH - 1) / BLOCKWIDTH
+ );
+ dim3 threads(BLOCKWIDTH);
+
+ AT_DISPATCH_FLOATING_TYPES(
+ vec.type(), "vecquant4matmul_batched_cuda", ([&] {
+ VecQuant4BatchMatMulColumnCompressionKernel<<>>(
+ vec.data(), mat.data(), mul.data(),
+ scales.data(), zeros.data(),
+ batch, heads, vec_row, height, width
+ );
+ })
+ );
+
+}
+
+template
+__global__ void VecQuant4BatchMatMulColumnCompressionKernel(
+ const scalar_t* __restrict__ vec,
+ const int* __restrict__ mat,
+ scalar_t* __restrict__ mul,
+ const scalar_t* __restrict__ scales,
+ const int* __restrict__ zeros,
+ int batch,
+ int heads,
+ int vec_row,
+ int height,
+ int width
+) {
+ int weight_total = batch * heads * height * width / 8;
+ int input_total = batch * heads * vec_row * height;
+ int out_total = batch * heads * vec_row * width;
+ int tid = threadIdx.x;
+ // h is index of height with step being BLOCKWIDTH
+ int h = BLOCKWIDTH * blockIdx.x;
+ // w is index of width with step being 1
+ int w = BLOCKWIDTH * blockIdx.y + tid;
+ if (w >= width && tid >= height) {
+ return;
+ }
+
+ __shared__ scalar_t blockvec[BLOCKWIDTH];
+ int k;
+ scalar_t w_tmp;
+
+ float weight[BLOCKWIDTH];
+
+ for (int b = 0; b < batch; ++b){
+ for (int head = 0; head < heads; ++head){
+ int batch_shift = b * heads + head;
+ for (k = 0; k < BLOCKWIDTH && h + k < height; ++k){
+ int i_w = (w / 8);
+ int w_bit = (w % 8) * 4;
+
+ int w_index = (batch_shift * height + h + k) * width / 8 + i_w;
+ if (w_index >= weight_total || w >= width) {
+ weight[k] = 0;
+ } else {
+ scalar_t scale = scales[batch_shift * height + h + k];
+ scalar_t zero = zeros[batch_shift * height + h + k];
+ w_tmp = ((as_unsigned(mat[w_index]) >> w_bit) & 0xF);
+ weight[k] = scale * (w_tmp - zero);
+ }
+ }
+
+ scalar_t res;
+ for (int vr = 0; vr < vec_row; ++vr){
+ res = 0;
+ int vec_index = (batch_shift * vec_row + vr) * height + blockIdx.x * BLOCKWIDTH + tid;
+ if (vec_index < input_total) {
+ blockvec[tid] = vec[vec_index];
+ } else {
+ blockvec[tid] = 0;
+ }
+
+ __syncthreads();
+ for (k = 0; k < BLOCKWIDTH && h + k < height; ++k){
+ // res is the dot product of BLOCKWIDTH elements (part of width)
+ res += weight[k] * blockvec[k];
+ }
+ // add res to the final result, final matrix shape: (batch, vec_row, width)
+ int out_index = (batch_shift * vec_row + vr) * width + w;
+ if (out_index < out_total) {
+ atomicAdd(&mul[out_index], res);
+ }
+ __syncthreads();
+ }
+ }
+ }
+}
+
+
+void vecquant8matmul_batched_old_cuda(
+ torch::Tensor vec,
+ torch::Tensor mat,
+ torch::Tensor mul,
+ torch::Tensor scales,
+ torch::Tensor zeros
+) {
+ int batch = vec.size(0);
+ int heads = vec.size(1);
+ int vec_row = vec.size(2);
+ int vec_height = vec.size(3);
+ int height = mat.size(2);
+ int width = mat.size(3);
+ int zero_width = zeros.size(2);
+
+ dim3 blocks(
+ (height + BLOCKWIDTH - 1) / BLOCKWIDTH,
+ (width + BLOCKWIDTH - 1) / BLOCKWIDTH
+ );
+ dim3 threads(BLOCKWIDTH);
+
+ AT_DISPATCH_FLOATING_TYPES(
+ vec.type(), "vecquant8matmul_batched_old_cuda", ([&] {
+ VecQuant8BatchMatMulKernel_old<<>>(
+ vec.data(), mat.data(), mul.data(),
+ scales.data(), zeros.data(),
+ batch, heads, vec_row, vec_height, height, width, zero_width
+ );
+ })
+ );
+}
+
+
+template
+__global__ void VecQuant8BatchMatMulKernel_old(
+ const scalar_t* __restrict__ vec,
+ const uint8_t* __restrict__ mat,
+ scalar_t* __restrict__ mul,
+ const scalar_t* __restrict__ scales,
+ const scalar_t* __restrict__ zeros,
+ int batch,
+ int heads,
+ int vec_row,
+ int vec_height,
+ int height,
+ int width,
+ int zero_width
+) {
+ int weight_total = batch * heads * height * width;
+ int input_total = batch * heads * vec_row * vec_height;
+ int out_total = batch * heads * vec_row * width;
+ int tid = threadIdx.x;
+ // h is index of height with step being BLOCKHEIGHT8
+ int h = BLOCKWIDTH * blockIdx.x;
+ // w is index of width with step being 1
+ int w = BLOCKWIDTH * blockIdx.y + tid;
+ if (w >= width && tid >= vec_height) {
+ return;
+ }
+
+ __shared__ scalar_t blockvec[BLOCKWIDTH];
+ // i is index of mat of block first row
+ int i = width * h + w;
+ int k;
+ scalar_t w_tmp;
+
+ float weight[BLOCKWIDTH];
+ for (int b = 0; b < batch; ++b){
+ for (int head = 0; head < heads; ++head){
+ int batch_shift = b * heads + head;
+ for (k = 0; k < BLOCKWIDTH && h + k < vec_height; ++k){
+ int k_w = k;
+ int w_index = batch_shift * height * width + i + (k_w * width);
+ if (w_index >= weight_total || w >= width) {
+ weight[k] = 0;
+ } else {
+ scalar_t scale = scales[batch_shift * width + w];
+ scalar_t zero = zeros[batch_shift * width + w];
+ w_tmp = as_unsigned(mat[w_index]);
+ weight[k] = scale * (w_tmp - zero);
+ }
+ }
+
+ scalar_t res;
+ for (int vr = 0; vr < vec_row; ++vr){
+ res = 0;
+ int vec_index = (batch_shift * vec_row + vr) * vec_height + blockIdx.x * BLOCKWIDTH + tid;
+ if (vec_index < input_total) {
+ blockvec[tid] = vec[vec_index];
+ } else {
+ blockvec[tid] = 0;
+ }
+
+ __syncthreads();
+ for (k = 0; k < BLOCKWIDTH && h + k < vec_height; ++k){
+ // res is the dot product of BLOCKWIDTH elements (part of width)
+ res += weight[k] * blockvec[k];
+ }
+ // add res to the final result, final matrix shape: (batch, vec_row, width)
+ int out_index = (batch_shift * vec_row + vr) * width + w;
+ if (out_index < out_total) {
+ atomicAdd(&mul[out_index], res);
+ }
+ __syncthreads();
+ }
+ }
+ }
+}
+
+
+
+void vecquant8matmul_batched_faster_cuda(
+ torch::Tensor vec,
+ torch::Tensor mat,
+ torch::Tensor mul,
+ torch::Tensor scales,
+ torch::Tensor zeros
+) {
+ int batch = vec.size(0);
+ int heads = vec.size(1);
+ int vec_row = vec.size(2);
+ int vec_height = vec.size(3);
+ int height = mat.size(2);
+ int width = mat.size(3);
+ int zero_width = zeros.size(2);
+
+ dim3 blocks(
+ (height + BLOCKWIDTH - 1) / BLOCKWIDTH,
+ (width + BLOCKWIDTH - 1) / BLOCKWIDTH
+ );
+ dim3 threads(BLOCKWIDTH);
+
+ VecQuant8BatchMatMulKernel_faster<<>>(
+ (half*) vec.data_ptr(),
+ (uint8_t*) mat.data_ptr(),
+ (half*) mul.data_ptr(),
+ (half*) scales.data_ptr(),
+ (half*) zeros.data_ptr(),
+ batch, heads, vec_row, vec_height, height, width, zero_width
+ );
+}
+
+
+
+__global__ void VecQuant8BatchMatMulKernel_faster(
+ const half* __restrict__ vec,
+ const uint8_t* __restrict__ mat,
+ half* __restrict__ mul,
+ const half* __restrict__ scales,
+ const half* __restrict__ zeros,
+ int batch,
+ int heads,
+ int vec_row,
+ int vec_height,
+ int height,
+ int width,
+ int zero_width
+) {
+ //int weight_total = batch * heads * height * width;
+ int input_total = batch * heads * vec_row * vec_height;
+ int out_total = batch * heads * vec_row * width;
+ int tid = threadIdx.x;
+ int h = BLOCKWIDTH * blockIdx.x;
+ int w = BLOCKWIDTH * blockIdx.y + tid;
+ if (w >= width && tid >= height) {
+ return;
+ }
+
+ __shared__ float blockvec[BLOCKWIDTH];
+ int i = width * h + w;
+ int k;
+ float w_tmp;
+
+ float weight[BLOCKWIDTH];
+ for (int b = 0; b < batch; ++b){
+ for (int head = 0; head < heads; ++head){
+ int batch_shift = b * heads + head;
+ for (k = 0; k < BLOCKWIDTH && h + k < vec_height; ++k){
+ int k_w = k;
+ int w_index = batch_shift * height * width + i + (k_w * width);
+ float scale = __half2float(scales[batch_shift * width + w]);
+ float zero = __half2float(zeros[batch_shift * width + w]);
+ w_tmp = as_unsigned(mat[w_index]);
+ weight[k] = scale *(w_tmp-zero);
+ }
+
+ float res;
+ for (int vr = 0; vr < vec_row; ++vr){
+ res = 0;
+ int vec_index = (batch_shift * vec_row + vr) * vec_height + blockIdx.x * BLOCKWIDTH + tid;
+ if (vec_index < input_total) {
+ blockvec[tid] = __half2float(vec[vec_index]);
+ } else {
+ blockvec[tid] = 0;
+ }
+ __syncthreads();
+ for (k = 0; k < BLOCKWIDTH && h + k < vec_height; ++k){
+ float temp_res = weight[k]*blockvec[k];
+ res += temp_res;
+ }
+ int out_index = (batch_shift * vec_row + vr) * width + w;
+ if (out_index < out_total) {
+ atomicAdd(&mul[out_index], __float2half(res));
+ }
+ __syncthreads();
+ }
+ }
+ }
+}
+
+
+
+
+void vecquant8matmul_batched_column_compression_faster_cuda(
+ torch::Tensor vec,
+ torch::Tensor mat,
+ torch::Tensor mul,
+ torch::Tensor scales,
+ torch::Tensor zeros
+) {
+ int batch = vec.size(0);
+ int heads = vec.size(1);
+ int vec_row = vec.size(2);
+ int height = vec.size(3);
+ int width = mat.size(3);
+
+ dim3 blocks(
+ (height + BLOCKWIDTH - 1) / BLOCKWIDTH,
+ (width + BLOCKWIDTH - 1) / BLOCKWIDTH
+ );
+ dim3 threads(BLOCKWIDTH);
+
+ VecQuant8BatchMatMulColumnCompressionKernel_faster<<>>(
+ (half*) vec.data_ptr(),
+ (uint8_t*) mat.data_ptr(),
+ (half*) mul.data_ptr(),
+ (half*) scales.data_ptr(),
+ (half*) zeros.data_ptr(),
+ batch, heads, vec_row, height, width
+ );
+
+}
+
+__global__ void VecQuant8BatchMatMulColumnCompressionKernel_faster(
+ const half* __restrict__ vec,
+ const uint8_t* __restrict__ mat,
+ half* __restrict__ mul,
+ const half* __restrict__ scales,
+ const half* __restrict__ zeros,
+ int batch,
+ int heads,
+ int vec_row,
+ int height,
+ int width
+) {
+ //int weight_total = batch * heads * height * width;
+ int input_total = batch * heads * vec_row * height;
+ int out_total = batch * heads * vec_row * width;
+ int tid = threadIdx.x;
+ int h = BLOCKWIDTH * blockIdx.x;
+ int w = BLOCKWIDTH * blockIdx.y + tid;
+ if (w >= width && tid >= height) {
+ return;
+ }
+
+ __shared__ float blockvec[BLOCKWIDTH];
+ int k;
+ float w_tmp;
+ float weight[BLOCKWIDTH];
+
+ for (int b = 0; b < batch; ++b){
+ for (int head = 0; head < heads; ++head){
+ int batch_shift = b * heads + head;
+ for (k = 0; k < BLOCKWIDTH; ++k){
+ int w_index = (batch_shift * height + h + k) * width + w;
+ float scale = __half2float(scales[batch_shift * height + h + k]);
+ float zero = __half2float(zeros[batch_shift * height + h + k]);
+ w_tmp = mat[w_index];
+ weight[k] = scale * (w_tmp-zero);
+ }
+
+ float res;
+ for (int vr = 0; vr < vec_row; ++vr){
+ res = 0;
+ int vec_index = (batch_shift * vec_row + vr) * height + blockIdx.x * BLOCKWIDTH + tid;
+ if (vec_index < input_total) {
+ blockvec[tid] = __half2float(vec[vec_index]);
+ } else {
+ blockvec[tid] = 0;
+ }
+ __syncthreads();
+ for (k = 0; k < BLOCKWIDTH; ++k){
+ res += weight[k]*blockvec[k];
+ }
+ int out_index = (batch_shift * vec_row + vr) * width + w;
+ if (out_index < out_total) {
+ atomicAdd(&mul[out_index], __float2half(res));
+ }
+ __syncthreads();
+ }
+ }
+ }
+}
+
+
+
+void vecquant8matmul_batched_column_compression_old_cuda(
+ torch::Tensor vec,
+ torch::Tensor mat,
+ torch::Tensor mul,
+ torch::Tensor scales,
+ torch::Tensor zeros
+) {
+ int batch = vec.size(0);
+ int heads = vec.size(1);
+ int vec_row = vec.size(2);
+ int height = vec.size(3);
+ int width = mat.size(3);
+
+ dim3 blocks(
+ (height + BLOCKWIDTH - 1) / BLOCKWIDTH,
+ (width + BLOCKWIDTH - 1) / BLOCKWIDTH
+ );
+ dim3 threads(BLOCKWIDTH);
+
+ AT_DISPATCH_FLOATING_TYPES(
+ vec.type(), "vecquant8matmul_batched_column_compression_old_cuda", ([&] {
+ VecQuant8BatchMatMulColumnCompressionKernel_old<<>>(
+ vec.data(), mat.data(), mul.data(),
+ scales.data(), zeros.data(),
+ batch, heads, vec_row, height, width
+ );
+ })
+ );
+
+}
+
+template
+__global__ void VecQuant8BatchMatMulColumnCompressionKernel_old(
+ const scalar_t* __restrict__ vec,
+ const uint8_t* __restrict__ mat,
+ scalar_t* __restrict__ mul,
+ const scalar_t* __restrict__ scales,
+ const scalar_t* __restrict__ zeros,
+ int batch,
+ int heads,
+ int vec_row,
+ int height,
+ int width
+) {
+ int weight_total = batch * heads * height * width;
+ int input_total = batch * heads * vec_row * height;
+ int out_total = batch * heads * vec_row * width;
+ int tid = threadIdx.x;
+ // h is index of height with step being BLOCKWIDTH
+ int h = BLOCKWIDTH * blockIdx.x;
+ // w is index of width with step being 1
+ int w = BLOCKWIDTH * blockIdx.y + tid;
+ if (w >= width && tid >= height) {
+ return;
+ }
+
+ __shared__ scalar_t blockvec[BLOCKWIDTH];
+ int k;
+ scalar_t w_tmp;
+
+ float weight[BLOCKWIDTH];
+
+ for (int b = 0; b < batch; ++b){
+ for (int head = 0; head < heads; ++head){
+ int batch_shift = b * heads + head;
+ for (k = 0; k < BLOCKWIDTH && h + k < height; ++k){
+ int w_index = (batch_shift * height + h + k) * width + w;
+ if (w_index >= weight_total || w >= width) {
+ weight[k] = 0;
+ } else {
+ scalar_t scale = scales[batch_shift * height + h + k];
+ scalar_t zero = zeros[batch_shift * height + h + k];
+ w_tmp = mat[w_index];
+ weight[k] = scale * (w_tmp - zero);
+ }
+ }
+
+ scalar_t res;
+ for (int vr = 0; vr < vec_row; ++vr){
+ res = 0;
+ int vec_index = (batch_shift * vec_row + vr) * height + blockIdx.x * BLOCKWIDTH + tid;
+ if (vec_index < input_total) {
+ blockvec[tid] = vec[vec_index];
+ } else {
+ blockvec[tid] = 0;
+ }
+
+ __syncthreads();
+ for (k = 0; k < BLOCKWIDTH && h + k < height; ++k){
+ // res is the dot product of BLOCKWIDTH elements (part of width)
+ res += weight[k] * blockvec[k];
+ }
+ // add res to the final result, final matrix shape: (batch, vec_row, width)
+ int out_index = (batch_shift * vec_row + vr) * width + w;
+ if (out_index < out_total) {
+ atomicAdd(&mul[out_index], res);
+ }
+ __syncthreads();
+ }
+ }
+ }
+}
+
+
+void vecquant4matmul_batched_old_cuda(
+ torch::Tensor vec,
+ torch::Tensor mat,
+ torch::Tensor mul,
+ torch::Tensor scales,
+ torch::Tensor zeros
+) {
+ int batch = vec.size(0);
+ int heads = vec.size(1);
+ int vec_row = vec.size(2);
+ int vec_height = vec.size(3);
+ int height = mat.size(2);
+ int width = mat.size(3);
+ int zero_width = zeros.size(2);
+
+ dim3 blocks(
+ (height + BLOCKHEIGHT_OLD4 - 1) / BLOCKHEIGHT_OLD4,
+ (width + BLOCKWIDTH - 1) / BLOCKWIDTH
+ );
+ dim3 threads(BLOCKWIDTH);
+
+ AT_DISPATCH_FLOATING_TYPES(
+ vec.type(), "vecquant4matmul_batched_old_cuda", ([&] {
+ VecQuant4BatchMatMulKernel_old<<>>(
+ vec.data(), mat.data(), mul.data(),
+ scales.data(), zeros.data(),
+ batch, heads, vec_row, vec_height, height, width, zero_width
+ );
+ })
+ );
+
+}
+
+template
+__global__ void VecQuant4BatchMatMulKernel_old(
+ const scalar_t* __restrict__ vec,
+ const uint8_t* __restrict__ mat,
+ scalar_t* __restrict__ mul,
+ const scalar_t* __restrict__ scales,
+ const scalar_t* __restrict__ zeros,
+ int batch,
+ int heads,
+ int vec_row,
+ int vec_height,
+ int height,
+ int width,
+ int zero_width
+) {
+ int weight_total = batch * heads * height * width;
+ int input_total = batch * heads * vec_row * vec_height;
+ int out_total = batch * heads * vec_row * width;
+ int tid = threadIdx.x;
+ // h is index of height with step being BLOCKHEIGHT_OLD4
+ int h = BLOCKHEIGHT_OLD4 * blockIdx.x;
+ // w is index of width with step being 1
+ int w = BLOCKWIDTH * blockIdx.y + tid;
+ if (w >= width && tid >= vec_height) {
+ return;
+ }
+
+ __shared__ scalar_t blockvec[BLOCKWIDTH];
+ // i is index of mat of block first row
+ int i = width * h + w;
+ int k;
+ scalar_t w_tmp;
+
+ float weight[BLOCKWIDTH];
+ for (int b = 0; b < batch; ++b){
+ for (int head = 0; head < heads; ++head){
+ int batch_shift = b * heads + head;
+ for (k = 0; k < BLOCKWIDTH && h*2 + k < vec_height; ++k){
+ int k_w = (k / 2);
+ int k_bit = (k % 2) * 4;
+ int w_index = batch_shift * height * width + i + (k_w * width);
+ if (w_index >= weight_total || w >= width) {
+ weight[k] = 0;
+ } else {
+ scalar_t scale = scales[batch_shift * width + w];
+ scalar_t zero = zeros[batch_shift * width + w];
+ w_tmp = ((as_unsigned(mat[w_index]) >> k_bit) & 0xF);
+ weight[k] = scale * (w_tmp - zero);
+ }
+ }
+
+ scalar_t res;
+ for (int vr = 0; vr < vec_row; ++vr){
+ res = 0;
+ int vec_index = (batch_shift * vec_row + vr) * vec_height + blockIdx.x * BLOCKWIDTH + tid;
+ if (vec_index < input_total) {
+ blockvec[tid] = vec[vec_index];
+ } else {
+ blockvec[tid] = 0;
+ }
+
+ __syncthreads();
+ for (k = 0; k < BLOCKWIDTH && h*2 + k < vec_height; ++k){
+ // res is the dot product of BLOCKWIDTH elements (part of width)
+ res += weight[k] * blockvec[k];
+ }
+ // add res to the final result, final matrix shape: (batch, vec_row, width)
+ int out_index = (batch_shift * vec_row + vr) * width + w;
+ if (out_index < out_total) {
+ atomicAdd(&mul[out_index], res);
+ }
+ __syncthreads();
+ }
+ }
+ }
+}
+
+
+
+
+
+void vecquant4matmul_batched_column_compression_old_cuda(
+ torch::Tensor vec,
+ torch::Tensor mat,
+ torch::Tensor mul,
+ torch::Tensor scales,
+ torch::Tensor zeros
+) {
+ int batch = vec.size(0);
+ int heads = vec.size(1);
+ int vec_row = vec.size(2);
+ int height = vec.size(3);
+ int width = mat.size(3);
+
+ dim3 blocks(
+ (height + BLOCKHEIGHT_OLD4 - 1) / BLOCKHEIGHT_OLD4,
+ (width + BLOCKWIDTH - 1) / BLOCKWIDTH
+ );
+ dim3 threads(BLOCKWIDTH);
+
+ AT_DISPATCH_FLOATING_TYPES(
+ vec.type(), "vecquant4matmul_batched_column_compression_old_cuda", ([&] {
+ VecQuant4BatchMatMulColumnCompressionKernel_old<<>>(
+ vec.data(), mat.data(), mul.data(),
+ scales.data(), zeros.data(),
+ batch, heads, vec_row, height, width
+ );
+ })
+ );
+
+}
+
+template
+__global__ void VecQuant4BatchMatMulColumnCompressionKernel_old(
+ const scalar_t* __restrict__ vec,
+ const uint8_t* __restrict__ mat,
+ scalar_t* __restrict__ mul,
+ const scalar_t* __restrict__ scales,
+ const scalar_t* __restrict__ zeros,
+ int batch,
+ int heads,
+ int vec_row,
+ int height,
+ int width
+) {
+ int weight_total = batch * heads * height * width;
+ int input_total = batch * heads * vec_row * height;
+ int out_total = batch * heads * vec_row * width;
+ int tid = threadIdx.x;
+ // h is index of height with step being BLOCKWIDTH
+ int h = BLOCKHEIGHT_OLD4 * blockIdx.x;
+ // w is index of width with step being 1
+ int w = BLOCKWIDTH * blockIdx.y + tid;
+ if (w >= width && tid >= height) {
+ return;
+ }
+
+ __shared__ scalar_t blockvec[BLOCKWIDTH];
+ int k;
+ scalar_t w_tmp;
+
+ float weight[BLOCKWIDTH];
+
+ for (int b = 0; b < batch; ++b){
+ for (int head = 0; head < heads; ++head){
+ int batch_shift = b * heads + head;
+ for (k = 0; k < BLOCKWIDTH && h*2 + k < height; ++k){
+ int k_w = (k / 2);
+ int k_bit = (k % 2) * 4;
+ int w_index = (batch_shift * height + h + k) * width + k_w;
+ if (w_index >= weight_total || w >= width) {
+ weight[k] = 0;
+ } else {
+ scalar_t scale = scales[batch_shift * height + h + k];
+ scalar_t zero = zeros[batch_shift * height + h + k];
+ w_tmp = ((as_unsigned(mat[w_index]) >> k_bit) & 0xF);
+ weight[k] = scale * (w_tmp - zero);
+ }
+ }
+
+ scalar_t res;
+ for (int vr = 0; vr < vec_row; ++vr){
+ res = 0;
+ int vec_index = (batch_shift * vec_row + vr) * height + blockIdx.x * BLOCKWIDTH + tid;
+ if (vec_index < input_total) {
+ blockvec[tid] = vec[vec_index];
+ } else {
+ blockvec[tid] = 0;
+ }
+
+ __syncthreads();
+ for (k = 0; k < BLOCKWIDTH && h*2 + k < height; ++k){
+ // res is the dot product of BLOCKWIDTH elements (part of width)
+ res += weight[k] * blockvec[k];
+ }
+ // add res to the final result, final matrix shape: (batch, vec_row, width)
+ int out_index = (batch_shift * vec_row + vr) * width + w;
+ if (out_index < out_total) {
+ atomicAdd(&mul[out_index], res);
+ }
+ __syncthreads();
+ }
+ }
+ }
+}
+
+
+
+
+
+void vecquant8matmul_batched_faster_old_cuda(
+ torch::Tensor vec,
+ torch::Tensor mat,
+ torch::Tensor mul,
+ torch::Tensor scales,
+ torch::Tensor zeros
+) {
+ int batch = vec.size(0);
+ int heads = vec.size(1);
+ int vec_row = vec.size(2);
+ int vec_height = vec.size(3);
+ int height = mat.size(2);
+ int width = mat.size(3);
+
+ dim3 blocks(
+ (height + BLOCKWIDTH - 1) / BLOCKWIDTH,
+ (width + BLOCKWIDTH - 1) / BLOCKWIDTH
+ );
+ dim3 threads(BLOCKWIDTH);
+
+ VecQuant8BatchMatMulKernel_faster_old<<>>(
+ (half*) vec.data_ptr(),
+ (uint8_t*) mat.data_ptr(),
+ (half*) mul.data_ptr(),
+ (half*) scales.data_ptr(),
+ (half*) zeros.data_ptr(),
+ batch, heads, vec_row, vec_height, height, width
+ );
+}
+
+
+__global__ void VecQuant8BatchMatMulKernel_faster_old(
+ const half* __restrict__ vec,
+ const uint8_t* __restrict__ mat,
+ half* __restrict__ mul,
+ const half* __restrict__ scales,
+ const half* __restrict__ zeros,
+ int batch,
+ int heads,
+ int vec_row,
+ int vec_height,
+ int height,
+ int width
+) {
+ int weight_total = batch * heads * height * width;
+ int input_total = batch * heads * vec_row * vec_height;
+ int out_total = batch * heads * vec_row * width;
+ int tid = threadIdx.x;
+ const int BLOCKWIDTH_half = BLOCKWIDTH/2;
+
+ int h = BLOCKWIDTH * blockIdx.x; //head_dim, dim=-1
+ int w = BLOCKWIDTH * blockIdx.y + tid; //seq-len, +0-256 ,dim=-2
+ /*
+ if (w >= width && tid >= vec_height) {
+ return;
+ }
+ */
+ __shared__ half blockvec[BLOCKWIDTH]; //256
+ int i = width * h + w;
+ int k;
+
+ half w_tmp1 = __float2half(0);
+ half w_tmp2 = __float2half(0);
+
+ half2 weight[BLOCKWIDTH_half];
+ for (int b = 0; b < batch; ++b){
+ for (int head = 0; head < heads; ++head){
+ int batch_shift = b * heads + head;
+ //int zero_index = batch_shift;
+ for (k = 0; k < BLOCKWIDTH_half; ++k){
+ int w_index1 = batch_shift * height * width + i + (2 * k * width); // [batch,head,h+k, w]
+ int w_index2 = batch_shift * height * width + i + ((2 * k + 1) * width);
+ int zero_index = batch_shift * width + w; // [batch,head, w]
+ if (w_index1 >= weight_total || w >= width || (2 * k + h) >= height) {
+ weight[k] = __float2half2_rn(0);
+ } else {
+ float zero_f=__half2float(zeros[zero_index]);
+ float scale_f= __half2float(scales[zero_index]);
+ if (w_index2 >= weight_total){
+ w_tmp1 = __float2half((as_unsigned(mat[w_index1]) -zero_f)*scale_f);
+ w_tmp2 = __float2half(0);
+ weight[k] = __halves2half2(w_tmp1,w_tmp2);
+ //printf("zero_index is %d w is %d height is %d width is %d w_index1 is %d w_tmp1 is %f w_tmp2 is %f zero is %f scale is %f low is %f high is %f \n ",zero_index,w,height, width,w_index1,__half2float(w_tmp1),__half2float(w_tmp2),zero_f,scale_f,__low2float(weight[k]),__high2float(weight[k]));
+ }else{
+ w_tmp1 = __int2half_rn(as_unsigned(mat[w_index1]));
+ w_tmp2 = __int2half_rn(as_unsigned(mat[w_index2]));
+
+ //weight[k] = __hmul2(__hsub2(__halves2half2(w_tmp1,w_tmp2), __halves2half2(zero,zero)),__halves2half2(scale,scale));
+ weight[k] = __hfma2(__halves2half2(w_tmp1,w_tmp2), __float2half2_rn(scale_f), __float2half2_rn(-(scale_f * zero_f)));
+ //printf("zero_index1 is %d zero_index2 is %d k is %d head is %d w is %d h is %d height is %d width is %d w_index1 is %d w_index2 is %d zero is %f scale is %f low is %f high is %f \n ",zero_index1,zero_index2,k,head,w,h,height, width,w_index1,w_index2,__half2float(zero1),__half2float(scale1),__low2float(weight[k]),__high2float(weight[k]));
+ }
+ }
+ }
+
+
+ for (int vr = 0; vr < vec_row; ++vr){
+ float res=0;
+ int vec_index = (batch_shift * vec_row + vr) * height + blockIdx.x * BLOCKWIDTH + tid;
+ int out_index = (batch_shift * vec_row + vr) * width + w;
+ if (vec_index < input_total) {
+ //blockvec[tid] = __half2float(vec[vec_index]);// [batch, head, vr, tid(seq_len dim+)]
+ blockvec[tid] = vec[vec_index];
+ //printf("width is %d height is %d h is %d w is %d vec_index is %d out_index is %d vec_row is %d vec_height is %d,vr is %d tid is %d blockvec is %f\n",width,height, h,w,vec_index,out_index,vec_row,vec_height,vr,tid,blockvec[tid]);
+ } else {
+ blockvec[tid] = __float2half(0);
+ }
+ __syncthreads();
+ if (out_index < out_total) {
+ for (k = 0; k < BLOCKWIDTH_half; ++k){
+ half2 res2 = __hmul2(weight[k],__halves2half2(blockvec[2*k],blockvec[2*k+1]));
+ res += __low2float(res2) + __high2float(res2);
+ }
+ atomicAdd(&mul[out_index], __float2half(res));
+ }
+ __syncthreads();
+ }
+ }
+ }
+}
+
+
+void vecquant8matmul_batched_column_compression_faster_old_cuda(
+ torch::Tensor vec, // [batch,heads, seq_q, seq_v]
+ torch::Tensor mat, // [batch,heads, seq_v, head_dim]
+ torch::Tensor mul, // [batch,heads, seq_q,head_dim]
+ torch::Tensor scales, // [batch,heads, head_dim]
+ torch::Tensor zeros
+) {
+ int batch = vec.size(0);
+ int heads = vec.size(1);
+ int vec_row = vec.size(2); //ql
+ int height = mat.size(2); //vl
+ int width = mat.size(3); //head_dim
+
+ dim3 blocks(
+ (height + BLOCKWIDTH - 1) / BLOCKWIDTH,
+ (width + BLOCKWIDTH - 1) / BLOCKWIDTH
+ );
+ dim3 threads(BLOCKWIDTH);
+
+ VecQuant8BatchMatMulColumnCompressionKernel_faster_old<<>>(
+ (half*) vec.data_ptr(),
+ (uint8_t*) mat.data_ptr(),
+ (half*) mul.data_ptr(),
+ (half*) scales.data_ptr(),
+ (half*) zeros.data_ptr(),
+ batch, heads, vec_row, height, width
+ );
+
+}
+
+
+__global__ void VecQuant8BatchMatMulColumnCompressionKernel_faster_old(
+ const half* __restrict__ vec, // [batch,heads, seq_q, seq_v]
+ const uint8_t* __restrict__ mat, // [batch,heads, seq_v, head_dim]
+ half* __restrict__ mul, // [batch,heads, seq_q,head_dim]
+ const half* __restrict__ scales, // [batch,heads, seq_v]
+ const half* __restrict__ zeros,
+ int batch,
+ int heads,
+ int vec_row, //seq_q
+ int height, //seq_v
+ int width //head_dim
+) {
+ int weight_total = batch * heads * height * width;
+ int input_total = batch * heads * vec_row * height;
+ int out_total = batch * heads * vec_row * width;
+ int tid = threadIdx.x;
+ int h = BLOCKWIDTH * blockIdx.x; // vl
+ int w = BLOCKWIDTH * blockIdx.y + tid; //head_dim + block
+ if (w >= width && tid >= height) {
+ return;
+ }
+ __shared__ half blockvec[BLOCKWIDTH];
+ int k;
+ half w_tmp1 = __float2half(0);
+ half w_tmp2 = __float2half(0);
+ int i = width * h + w;
+ const int BLOCKWIDTH_half = BLOCKWIDTH/2;
+ half2 weight[BLOCKWIDTH_half];
+
+ for (int b = 0; b < batch; ++b){
+ for (int head = 0; head < heads; ++head){
+ int batch_shift = b * heads + head;
+ //int zero_index = batch_shift;
+ for (k = 0; k < BLOCKWIDTH_half; ++k){
+ int w_index1 = batch_shift * height * width + i + (2 * k) * width; // [batch,head, h+k, w]
+ int w_index2 = batch_shift * height * width + i + ((2 * k + 1) * width);
+ int zero_index1 = batch_shift * height + h + 2*k; // [batch,head, w]
+ int zero_index2 = batch_shift * height + h + 2*k+1; // [batch,head, w]
+
+ if (w_index1 >= weight_total || (2 * k + h)>=height) {
+ weight[k]=__float2half2_rn(0);
+ } else{
+ //int zero_index = batch_shift + h; // [batch,head, w]
+ //float scale_f1 = __half2float(scales[zero_index1]);
+ //float zero_f1 = __half2float(zeros[zero_index1]);
+ if (w_index2>=weight_total){
+ w_tmp1 = __float2half((as_unsigned(mat[w_index1]) - __half2float(zeros[zero_index1]))* __half2float(scales[zero_index1]));
+ w_tmp2 = __float2half(0);
+ weight[k] = __halves2half2(w_tmp1,w_tmp2);
+ //printf("zero_index is %d k is %d w is %d head is %d height is %d width is %d w_index1 is %d w_tmp1 is %f w_tmp2 is %f zero is %f scale is %f low is %f high is %f \n ",zero_index,k,w,head,height, width,w_index1,__half2float(w_tmp1),__half2float(w_tmp2),zero_f,scale_f,__low2float(weight[k]),__high2float(weight[k]));
+ }else{
+ w_tmp1 = __int2half_rn(as_unsigned(mat[w_index1]));
+ w_tmp2 = __int2half_rn(as_unsigned(mat[w_index2]));
+ half zero1=zeros[zero_index1];
+ half zero2=zeros[zero_index2];
+ half scale1=scales[zero_index1];
+ half scale2=scales[zero_index2];
+ weight[k] = __hmul2(__hsub2(__halves2half2(w_tmp1,w_tmp2), __halves2half2(zero1,zero2)),__halves2half2(scale1,scale2));
+ //weight[k] = __hfma2(__halves2half2(w_tmp1,w_tmp2), __float2half2_rn(scale_f), __float2half2_rn(-(scale_f * zero_f)));
+ //printf("zero_index1 is %d zero_index2 is %d k is %d head is %d w is %d h is %d height is %d width is %d w_index1 is %d w_index2 is %d zero is %f scale is %f low is %f high is %f \n ",zero_index1,zero_index2,k,head,w,h,height, width,w_index1,w_index2,__half2float(zero1),__half2float(scale1),__low2float(weight[k]),__high2float(weight[k]));
+ }
+ }
+ }
+
+
+ for (int vr = 0; vr < vec_row; ++vr){
+ float res=0;
+ int vec_index = (batch_shift * vec_row + vr) * height + blockIdx.x * BLOCKWIDTH + tid;
+ int out_index = (batch_shift * vec_row + vr) * width + w;
+
+ if (vec_index < input_total) {
+ //blockvec[tid] = __half2float(vec[vec_index]);
+ blockvec[tid] = vec[vec_index];
+ //printf("vec_index is %d out_index is %d vec_row is %d ,vr is %d tid is %d blockvec is %f\n",vec_index,out_index,vec_row,vr,tid,blockvec[tid]);
+ } else {
+ blockvec[tid] = __float2half(0);
+ //blockvec[tid] = 0;
+ }
+ __syncthreads();
+ if (out_index < out_total) {
+ for (k = 0; k < BLOCKWIDTH_half; ++k){
+ half2 res2 = __hmul2(weight[k],__halves2half2(blockvec[2*k],blockvec[2*k+1]));
+ res += __low2float(res2) + __high2float(res2);
+ }
+ atomicAdd(&mul[out_index], __float2half(res));
+ }
+ __syncthreads();
+ }
+ }
+ }
+}
diff --git a/mft_peft_hf/src/model/qwen/configuration_qwen.py b/mftcoder_accelerate/src/model/qwen/configuration_qwen.py
similarity index 53%
rename from mft_peft_hf/src/model/qwen/configuration_qwen.py
rename to mftcoder_accelerate/src/model/qwen/configuration_qwen.py
index 502ef6c..f8fe2cb 100644
--- a/mft_peft_hf/src/model/qwen/configuration_qwen.py
+++ b/mftcoder_accelerate/src/model/qwen/configuration_qwen.py
@@ -9,61 +9,49 @@
class QWenConfig(PretrainedConfig):
model_type = "qwen"
keys_to_ignore_at_inference = ["past_key_values"]
- attribute_map = {
- "hidden_size": "n_embd",
- "num_attention_heads": "n_head",
- "max_position_embeddings": "n_positions",
- "num_hidden_layers": "n_layer",
- }
def __init__(
self,
- vocab_size=151851,
- n_embd=4096,
- n_layer=32,
- n_head=32,
- n_inner=None,
- embd_pdrop=0.0,
- attn_pdrop=0.0,
- layer_norm_epsilon=1e-5,
+ vocab_size=151936,
+ hidden_size=4096,
+ num_hidden_layers=32,
+ num_attention_heads=32,
+ emb_dropout_prob=0.0,
+ attn_dropout_prob=0.0,
+ layer_norm_epsilon=1e-6,
initializer_range=0.02,
+ max_position_embeddings=8192,
scale_attn_weights=True,
use_cache=True,
- eos_token_id=151643,
- apply_residual_connection_post_layernorm=False,
bf16=False,
fp16=False,
fp32=False,
kv_channels=128,
rotary_pct=1.0,
rotary_emb_base=10000,
- use_dynamic_ntk=False,
- use_logn_attn=False,
- use_flash_attn=True,
- ffn_hidden_size=22016,
+ use_dynamic_ntk=True,
+ use_logn_attn=True,
+ use_flash_attn="auto",
+ intermediate_size=22016,
no_bias=True,
tie_word_embeddings=False,
+ use_cache_quantization=False,
+ use_cache_kernel=False,
+ softmax_in_fp32=False,
**kwargs,
):
- self.eos_token_id = eos_token_id
- super().__init__(
- eos_token_id=eos_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs
- )
-
self.vocab_size = vocab_size
- self.n_embd = n_embd
- self.n_layer = n_layer
- self.n_head = n_head
- self.n_inner = n_inner
- self.embd_pdrop = embd_pdrop
- self.attn_pdrop = attn_pdrop
+ self.hidden_size = hidden_size
+ self.intermediate_size = intermediate_size
+ self.num_hidden_layers = num_hidden_layers
+ self.num_attention_heads = num_attention_heads
+ self.emb_dropout_prob = emb_dropout_prob
+ self.attn_dropout_prob = attn_dropout_prob
self.layer_norm_epsilon = layer_norm_epsilon
self.initializer_range = initializer_range
self.scale_attn_weights = scale_attn_weights
self.use_cache = use_cache
- self.apply_residual_connection_post_layernorm = (
- apply_residual_connection_post_layernorm
- )
+ self.max_position_embeddings = max_position_embeddings
self.bf16 = bf16
self.fp16 = fp16
self.fp32 = fp32
@@ -73,6 +61,11 @@ def __init__(
self.use_dynamic_ntk = use_dynamic_ntk
self.use_logn_attn = use_logn_attn
self.use_flash_attn = use_flash_attn
- self.ffn_hidden_size = ffn_hidden_size
self.no_bias = no_bias
- self.tie_word_embeddings = tie_word_embeddings
+ self.use_cache_quantization = use_cache_quantization
+ self.use_cache_kernel = use_cache_kernel
+ self.softmax_in_fp32 = softmax_in_fp32
+ super().__init__(
+ tie_word_embeddings=tie_word_embeddings,
+ **kwargs
+ )
diff --git a/mftcoder_accelerate/src/model/qwen/cpp_kernels.py b/mftcoder_accelerate/src/model/qwen/cpp_kernels.py
new file mode 100644
index 0000000..d9cee70
--- /dev/null
+++ b/mftcoder_accelerate/src/model/qwen/cpp_kernels.py
@@ -0,0 +1,55 @@
+from torch.utils import cpp_extension
+import pathlib
+import os
+import subprocess
+
+def _get_cuda_bare_metal_version(cuda_dir):
+ raw_output = subprocess.check_output([cuda_dir + "/bin/nvcc", "-V"],
+ universal_newlines=True)
+ output = raw_output.split()
+ release_idx = output.index("release") + 1
+ release = output[release_idx].split(".")
+ bare_metal_major = release[0]
+ bare_metal_minor = release[1][0]
+
+ return raw_output, bare_metal_major, bare_metal_minor
+
+def _create_build_dir(buildpath):
+ try:
+ os.mkdir(buildpath)
+ except OSError:
+ if not os.path.isdir(buildpath):
+ print(f"Creation of the build directory {buildpath} failed")
+
+# Check if cuda 11 is installed for compute capability 8.0
+cc_flag = []
+_, bare_metal_major, bare_metal_minor = _get_cuda_bare_metal_version(cpp_extension.CUDA_HOME)
+if int(bare_metal_major) >= 11:
+ cc_flag.append('-gencode')
+ cc_flag.append('arch=compute_80,code=sm_80')
+ if int(bare_metal_minor) >= 7:
+ cc_flag.append('-gencode')
+ cc_flag.append('arch=compute_90,code=sm_90')
+
+# Build path
+srcpath = pathlib.Path(__file__).parent.absolute()
+buildpath = srcpath / 'build'
+_create_build_dir(buildpath)
+
+def _cpp_extention_load_helper(name, sources, extra_cuda_flags):
+ return cpp_extension.load(
+ name=name,
+ sources=sources,
+ build_directory=buildpath,
+ extra_cflags=['-O3', ],
+ extra_cuda_cflags=['-O3',
+ '-gencode', 'arch=compute_70,code=sm_70',
+ '--use_fast_math'] + extra_cuda_flags + cc_flag,
+ verbose=1
+ )
+
+extra_flags = []
+
+cache_autogptq_cuda_256_sources = ["./cache_autogptq_cuda_256.cpp",
+ "./cache_autogptq_cuda_kernel_256.cu"]
+cache_autogptq_cuda_256 = _cpp_extention_load_helper("cache_autogptq_cuda_256", cache_autogptq_cuda_256_sources, extra_flags)
diff --git a/mft_peft_hf/src/model/qwen/modeling_qwen.py b/mftcoder_accelerate/src/model/qwen/modeling_qwen.py
similarity index 62%
rename from mft_peft_hf/src/model/qwen/modeling_qwen.py
rename to mftcoder_accelerate/src/model/qwen/modeling_qwen.py
index 238fe91..45c0d16 100644
--- a/mft_peft_hf/src/model/qwen/modeling_qwen.py
+++ b/mftcoder_accelerate/src/model/qwen/modeling_qwen.py
@@ -3,14 +3,16 @@
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
+import copy
import importlib
import math
+import pathlib
from typing import TYPE_CHECKING, Optional, Tuple, Union, Callable, List, Any, Generator
import torch
import torch.nn.functional as F
import torch.utils.checkpoint
-from torch.cuda.amp import autocast
+import warnings
from torch.nn import CrossEntropyLoss
from transformers import PreTrainedTokenizer, GenerationConfig, StoppingCriteriaList
@@ -35,6 +37,8 @@
SUPPORT_CUDA = torch.cuda.is_available()
SUPPORT_BF16 = SUPPORT_CUDA and torch.cuda.is_bf16_supported()
SUPPORT_FP16 = SUPPORT_CUDA and torch.cuda.get_device_capability(0)[0] >= 7
+SUPPORT_TORCH2 = hasattr(torch, '__version__') and int(torch.__version__.split(".")[0]) >= 2
+
from .configuration_qwen import QWenConfig
from .qwen_generation_utils import (
@@ -66,13 +70,18 @@
向model.chat()传入参数stream的用法可能存在Bug，该用法已被废弃，将在未来被移除。请使用model.chat_stream(...)代替model.chat(..., stream=True)。
"""
+_ERROR_INPUT_CPU_QUERY_WITH_FLASH_ATTN_ACTIVATED = """\
+We detect you have activated flash attention support, but running model computation on CPU. Please make sure that your input data has been placed on GPU. If you actually want to run CPU computation, please following the readme and set device_map="cpu" to disable flash attention when loading the model (calling AutoModelForCausalLM.from_pretrained).
+检测到您的模型已激活了flash attention支持，但正在执行CPU运算任务。如使用flash attention，请您确认模型输入已经传到GPU上。如果您确认要执行CPU运算，请您在载入模型（调用AutoModelForCausalLM.from_pretrained）时，按照readme说法，指定device_map="cpu"以禁用flash attention。
+"""
+
apply_rotary_emb_func = None
rms_norm = None
flash_attn_unpadded_func = None
-
+flash_attn_func = None
def _import_flash_attn():
- global apply_rotary_emb_func, rms_norm, flash_attn_unpadded_func
+ global apply_rotary_emb_func, rms_norm, flash_attn_unpadded_func, flash_attn_func
try:
from flash_attn.layers.rotary import apply_rotary_emb_func as __apply_rotary_emb_func
apply_rotary_emb_func = __apply_rotary_emb_func
@@ -93,20 +102,49 @@ def _import_flash_attn():
try:
import flash_attn
+ _flash_attn_func = None
if not hasattr(flash_attn, '__version__'):
from flash_attn.flash_attn_interface import flash_attn_unpadded_func as __flash_attn_unpadded_func
else:
if int(flash_attn.__version__.split(".")[0]) >= 2:
+ if int(flash_attn.__version__.split(".")[1]) >= 1:
+ from flash_attn.flash_attn_interface import flash_attn_func as _flash_attn_func
from flash_attn.flash_attn_interface import flash_attn_varlen_func as __flash_attn_unpadded_func
else:
from flash_attn.flash_attn_interface import flash_attn_unpadded_func as __flash_attn_unpadded_func
flash_attn_unpadded_func = __flash_attn_unpadded_func
+ flash_attn_func = _flash_attn_func
except ImportError:
logger.warn(
"Warning: import flash_attn fail, please install FlashAttention to get higher efficiency "
"https://github.com/Dao-AILab/flash-attention"
)
+def quantize_cache_v(fdata, bits, qmax, qmin):
+ # b, s, head, h-dim->b, head, s, h-dim
+ qtype = torch.uint8
+ device = fdata.device
+ shape = fdata.shape
+
+ fdata_cal = torch.flatten(fdata, 2)
+ fmax = torch.amax(fdata_cal, dim=-1, keepdim=True)
+ fmin = torch.amin(fdata_cal, dim=-1, keepdim=True)
+ # Compute params
+ if qmax.device != fmax.device:
+ qmax = qmax.to(device)
+ qmin = qmin.to(device)
+ scale = (fmax - fmin) / (qmax - qmin)
+ zero = qmin - fmin / scale
+ scale = scale.unsqueeze(-1).repeat(1,1,shape[2],1).contiguous()
+ zero = zero.unsqueeze(-1).repeat(1,1,shape[2],1).contiguous()
+ # Quantize
+ res_data = fdata / scale + zero
+ qdata = torch.clamp(res_data, qmin, qmax).to(qtype)
+ return qdata.contiguous(), scale, zero
+
+def dequantize_cache_torch(qdata, scale, zero):
+ data = scale * (qdata - zero)
+ return data
class FlashSelfAttention(torch.nn.Module):
def __init__(
@@ -126,11 +164,33 @@ def __init__(
self.softmax_scale = softmax_scale
self.dropout_p = attention_dropout
- def forward(self, q, k, v):
+ def unpad_input(self, hidden_states, attention_mask):
+ valid_mask = attention_mask.squeeze(1).squeeze(1).eq(0)
+ seqlens_in_batch = valid_mask.sum(dim=-1, dtype=torch.int32)
+ indices = torch.nonzero(valid_mask.flatten(), as_tuple=False).flatten()
+ max_seqlen_in_batch = seqlens_in_batch.max().item()
+ cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0))
+ hidden_states = hidden_states[indices]
+ return hidden_states, indices, cu_seqlens, max_seqlen_in_batch
+
+ def pad_input(self, hidden_states, indices, batch, seqlen):
+ output = torch.zeros(batch * seqlen, *hidden_states.shape[1:], device=hidden_states.device,
+ dtype=hidden_states.dtype)
+ output[indices] = hidden_states
+ return rearrange(output, '(b s) ... -> b s ...', b=batch)
+
+ def forward(self, q, k, v, attention_mask=None):
assert all((i.dtype in [torch.float16, torch.bfloat16] for i in (q, k, v)))
assert all((i.is_cuda for i in (q, k, v)))
batch_size, seqlen_q = q.shape[0], q.shape[1]
seqlen_k = k.shape[1]
+ seqlen_out = seqlen_q
+
+ if flash_attn_func is not None and batch_size == 1:
+ dropout_p = self.dropout_p if self.training else 0
+ output = flash_attn_func(q, k, v, dropout_p, softmax_scale=self.softmax_scale, causal=self.causal)
+ return output
+
q, k, v = [rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v]]
cu_seqlens_q = torch.arange(
0,
@@ -140,13 +200,14 @@ def forward(self, q, k, v):
device=q.device,
)
- if self.training:
- assert seqlen_k == seqlen_q
-
- is_causal = self.causal
- cu_seqlens_k = cu_seqlens_q
+ if batch_size > 1 and attention_mask is not None:
+ k, indices_k, cu_seqlens_k, seqlen_k = self.unpad_input(k, attention_mask)
+ if q.size(0) == v.size(0):
+ q = q[indices_k]
+ cu_seqlens_q = cu_seqlens_k
+ seqlen_q = seqlen_k
+ v = v[indices_k]
else:
- is_causal = seqlen_q == seqlen_k
cu_seqlens_k = torch.arange(
0,
(batch_size + 1) * seqlen_k,
@@ -154,7 +215,15 @@ def forward(self, q, k, v):
dtype=torch.int32,
device=q.device,
)
- self.dropout_p = 0
+
+ if self.training:
+ assert seqlen_k == seqlen_q
+ is_causal = self.causal
+ dropout_p = self.dropout_p
+ else:
+ is_causal = seqlen_q == seqlen_k
+ dropout_p = 0
+
output = flash_attn_unpadded_func(
q,
k,
@@ -163,30 +232,23 @@ def forward(self, q, k, v):
cu_seqlens_k,
seqlen_q,
seqlen_k,
- self.dropout_p,
+ dropout_p,
softmax_scale=self.softmax_scale,
causal=is_causal,
)
-
- output = rearrange(output, "(b s) ... -> b s ...", b=batch_size)
+ if batch_size > 1 and attention_mask is not None and seqlen_q == seqlen_k:
+ output = self.pad_input(output, indices_k, batch_size, seqlen_out)
+ else:
+ new_shape = (batch_size, output.shape[0] // batch_size) + output.shape[1:]
+ output = output.view(new_shape)
return output
class QWenAttention(nn.Module):
- def __init__(self, config, layer_number=None):
+ def __init__(self, config):
super().__init__()
- max_positions = config.max_position_embeddings
- self.register_buffer(
- "bias",
- torch.tril(
- torch.ones((max_positions, max_positions), dtype=torch.bool)
- ).view(1, 1, max_positions, max_positions),
- persistent=False,
- )
self.register_buffer("masked_bias", torch.tensor(-1e4), persistent=False)
- self.layer_number = max(1, layer_number)
- self.params_dtype = config.params_dtype
self.seq_length = config.seq_length
self.hidden_size = config.hidden_size
@@ -197,8 +259,6 @@ def __init__(self, config, layer_number=None):
self.use_flash_attn = config.use_flash_attn
self.scale_attn_weights = True
- self.layer_idx = None
-
self.projection_size = config.kv_channels * config.num_attention_heads
assert self.projection_size % config.num_attention_heads == 0
@@ -219,25 +279,10 @@ def __init__(self, config, layer_number=None):
and not self.is_fp32
):
self.core_attention_flash = FlashSelfAttention(
- causal=True, attention_dropout=config.attn_pdrop
+ causal=True, attention_dropout=config.attn_dropout_prob
)
-
self.bf16 = config.bf16
- if config.rotary_pct == 1.0:
- self.rotary_ndims = None
- else:
- assert config.rotary_pct < 1
- self.rotary_ndims = int(
- self.hidden_size_per_attention_head * config.rotary_pct
- )
- dim = (
- self.rotary_ndims
- if self.rotary_ndims is not None
- else self.hidden_size_per_attention_head
- )
- self.rotary_emb = RotaryEmbedding(dim, base=config.rotary_emb_base)
-
self.use_dynamic_ntk = config.use_dynamic_ntk
self.use_logn_attn = config.use_logn_attn
@@ -245,104 +290,104 @@ def __init__(self, config, layer_number=None):
math.log(i, self.seq_length) if i > self.seq_length else 1
for i in range(1, 32768)
]
- self.logn_tensor = torch.Tensor(logn_list)[None, :, None, None]
- self._ntk_cached = 1.0
-
- self.attn_dropout = nn.Dropout(config.attn_pdrop)
-
- def _attn(self, query, key, value, attention_mask=None, head_mask=None):
- attn_weights = torch.matmul(query, key.transpose(-1, -2))
+ logn_tensor = torch.tensor(logn_list)[None, :, None, None]
+ self.register_buffer("logn_tensor", logn_tensor, persistent=False)
+
+ self.attn_dropout = nn.Dropout(config.attn_dropout_prob)
+ self.softmax_in_fp32 = config.softmax_in_fp32 if hasattr(config, 'softmax_in_fp32') else False
+ self.use_cache_quantization = config.use_cache_quantization if hasattr(config, 'use_cache_quantization') else False
+ self.use_cache_kernel = config.use_cache_kernel if hasattr(config,'use_cache_kernel') else False
+ cache_dtype = torch.float
+ if self.bf16:
+ cache_dtype=torch.bfloat16
+ elif config.fp16:
+ cache_dtype = torch.float16
+ self.cache_qmax = torch.tensor(torch.iinfo(torch.uint8).max, dtype=cache_dtype)
+ self.cache_qmin = torch.tensor(torch.iinfo(torch.uint8).min, dtype=cache_dtype)
+
+ if config.use_cache_quantization and config.use_cache_kernel:
+ # pre check if the support files existing
+ module_root = pathlib.Path(__file__).parent
+ src_files = ("cache_autogptq_cuda_256.cpp", "cache_autogptq_cuda_kernel_256.cu")
+ if any(not (module_root/src).is_file() for src in src_files):
+ warnings.warn("KV cache kernel source files (.cpp and .cu) not found.")
+ self.cache_kernels = None
+ else:
+ try:
+ from .cpp_kernels import cache_autogptq_cuda_256
+ self.cache_kernels = cache_autogptq_cuda_256
+ except ImportError:
+ warnings.warn("Failed to import KV cache kernels.")
+ self.cache_kernels = None
+
+ def _attn(self, query, key, value, causal_mask=None, attention_mask=None, head_mask=None):
+ device = query.device
+ if self.use_cache_quantization:
+ qk, qk_scale, qk_zero = key
+ if self.use_cache_kernel and self.cache_kernels is not None:
+ shape = query.shape[:-1] + (qk.shape[-2],)
+ attn_weights = torch.zeros(shape, dtype=torch.float16, device=device)
+ self.cache_kernels.vecquant8matmul_batched_faster_old(
+ query.contiguous() if query.dtype == torch.float16 else query.to(torch.float16).contiguous(),
+ qk.transpose(-1, -2).contiguous(),
+ attn_weights,
+ qk_scale.contiguous() if qk_scale.dtype == torch.float16 else qk_scale.to(torch.float16).contiguous(),
+ qk_zero.contiguous()if qk_zero.dtype == torch.float16 else qk_zero.to(torch.float16).contiguous())
+ # attn_weights = attn_weights.to(query.dtype).contiguous()
+ else:
+ key = dequantize_cache_torch(qk, qk_scale, qk_zero)
+ attn_weights = torch.matmul(query, key.transpose(-1, -2))
+ else:
+ attn_weights = torch.matmul(query, key.transpose(-1, -2))
if self.scale_attn_weights:
- attn_weights = attn_weights / torch.full(
- [],
- value.size(-1) ** 0.5,
- dtype=attn_weights.dtype,
- device=attn_weights.device,
- )
+ if self.use_cache_quantization:
+ size_temp = value[0].size(-1)
+ else:
+ size_temp = value.size(-1)
+ attn_weights = attn_weights / (size_temp ** 0.5)
- query_length, key_length = query.size(-2), key.size(-2)
- causal_mask = self.bias[
- :, :, key_length - query_length : key_length, :key_length
- ]
mask_value = torch.finfo(attn_weights.dtype).min
- mask_value = torch.full([], mask_value, dtype=attn_weights.dtype).to(
- attn_weights.device
- )
- attn_weights = torch.where(
- causal_mask, attn_weights.to(attn_weights.dtype), mask_value
- )
-
- if attention_mask is not None:
- # Apply the attention mask
- attn_weights = attn_weights + attention_mask
-
- attn_weights = nn.functional.softmax(attn_weights, dim=-1)
-
- attn_weights = attn_weights.type(value.dtype)
- attn_weights = self.attn_dropout(attn_weights)
-
- if head_mask is not None:
- attn_weights = attn_weights * head_mask
-
- attn_output = torch.matmul(attn_weights, value)
- attn_output = attn_output.transpose(1, 2)
-
- return attn_output, attn_weights
-
- def _upcast_and_reordered_attn(
- self, query, key, value, attention_mask=None, head_mask=None
- ):
- bsz, num_heads, q_seq_len, dk = query.size()
- _, _, k_seq_len, _ = key.size()
-
- attn_weights = torch.empty(
- bsz * num_heads,
- q_seq_len,
- k_seq_len,
- dtype=torch.float32,
- device=query.device,
- )
-
- scale_factor = 1.0
- if self.scale_attn_weights:
- scale_factor /= float(value.size(-1)) ** 0.5
-
- with autocast(enabled=False):
- q, k = query.reshape(-1, q_seq_len, dk), key.transpose(-1, -2).reshape(
- -1, dk, k_seq_len
- )
- attn_weights = torch.baddbmm(
- attn_weights, q.float(), k.float(), beta=0, alpha=scale_factor
+ if causal_mask is not None:
+ attn_weights = torch.where(
+ causal_mask, attn_weights.to(attn_weights.dtype), mask_value
)
- attn_weights = attn_weights.reshape(bsz, num_heads, q_seq_len, k_seq_len)
-
- query_length, key_length = query.size(-2), key.size(-2)
- causal_mask = self.bias[
- :, :, key_length - query_length : key_length, :key_length
- ]
- mask_value = torch.finfo(attn_weights.dtype).min
- mask_value = torch.tensor(mask_value, dtype=attn_weights.dtype).to(
- attn_weights.device
- )
- attn_weights = torch.where(causal_mask, attn_weights, mask_value)
if attention_mask is not None:
attn_weights = attn_weights + attention_mask
- attn_weights = nn.functional.softmax(attn_weights, dim=-1)
+ if self.softmax_in_fp32:
+ attn_weights = nn.functional.softmax(attn_weights.float(), dim=-1)
+ else:
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1)
- if attn_weights.dtype != torch.float32:
- raise RuntimeError(
- "Error with upcasting, attn_weights does not have dtype torch.float32"
- )
- attn_weights = attn_weights.type(value.dtype)
+ attn_weights = attn_weights.type(query.dtype)
attn_weights = self.attn_dropout(attn_weights)
if head_mask is not None:
attn_weights = attn_weights * head_mask
- attn_output = torch.matmul(attn_weights, value)
+ if self.use_cache_quantization:
+ qv, qv_scale, qv_zero = value
+ if self.use_cache_kernel and self.cache_kernels is not None:
+ shape = attn_weights.shape[:-1] + (query.shape[-1],)
+ attn_output = torch.zeros(shape, dtype=torch.float16, device=device)
+ self.cache_kernels.vecquant8matmul_batched_column_compression_faster_old(
+ attn_weights.contiguous() if attn_weights.dtype == torch.float16 else attn_weights.to(torch.float16).contiguous(),
+ qv.contiguous(), # dtype: int32
+ attn_output,
+ qv_scale.contiguous() if qv_scale.dtype == torch.float16 else qv_scale.to(torch.float16).contiguous(),
+ qv_zero.contiguous() if qv_zero.dtype == torch.float16 else qv_zero.to(torch.float16).contiguous())
+ if attn_output.dtype != query.dtype:
+ attn_output = attn_output.to(query.dtype)
+ attn_weights = attn_weights.to(query.dtype)
+ else:
+ value = dequantize_cache_torch(qv, qv_scale, qv_zero)
+ attn_output = torch.matmul(attn_weights, value)
+ else:
+ attn_output = torch.matmul(attn_weights, value)
+
+ attn_output = attn_output.transpose(1, 2)
return attn_output, attn_weights
@@ -359,6 +404,7 @@ def _merge_heads(self, tensor, num_heads, attn_head_size):
def forward(
self,
hidden_states: Optional[Tuple[torch.FloatTensor]],
+ rotary_pos_emb_list: Optional[List[List[torch.Tensor]]] = None,
layer_past: Optional[Tuple[torch.Tensor]] = None,
attention_mask: Optional[torch.FloatTensor] = None,
head_mask: Optional[torch.FloatTensor] = None,
@@ -367,64 +413,80 @@ def forward(
output_attentions: Optional[bool] = False,
use_cache: Optional[bool] = False,
):
-
mixed_x_layer = self.c_attn(hidden_states)
+
query, key, value = mixed_x_layer.split(self.split_size, dim=2)
query = self._split_heads(query, self.num_heads, self.head_dim)
key = self._split_heads(key, self.num_heads, self.head_dim)
value = self._split_heads(value, self.num_heads, self.head_dim)
- kv_seq_len = hidden_states.size()[1]
- if layer_past:
- # layer past[0] shape: bs * seq_len * head_num * dim
- kv_seq_len += layer_past[0].shape[1]
- if (
- self.use_dynamic_ntk
- and kv_seq_len == hidden_states.size()[1]
- and not self.training
- ):
- context_value = math.log(kv_seq_len / self.seq_length, 2) + 1
- ntk_alpha = 2 ** math.ceil(context_value) - 1
- ntk_alpha = max(ntk_alpha, 1)
- self._ntk_cached = ntk_alpha
- else:
- ntk_alpha = self._ntk_cached
- rotary_pos_emb = self.rotary_emb(kv_seq_len, ntk_alpha=ntk_alpha).to(
- hidden_states.device
- )
-
- if rotary_pos_emb is not None:
- if isinstance(rotary_pos_emb, tuple):
- rotary_pos_emb = rotary_pos_emb
- else:
+ if rotary_pos_emb_list is not None:
+ cur_len = query.shape[1]
+ if len(rotary_pos_emb_list) == 1:
+ rotary_pos_emb = rotary_pos_emb_list[0]
+ rotary_pos_emb = [i[:, -cur_len:, :, :] for i in rotary_pos_emb]
rotary_pos_emb = (rotary_pos_emb,) * 2
+ q_pos_emb, k_pos_emb = rotary_pos_emb
+ # Slice the pos emb for current inference
+ query = apply_rotary_pos_emb(query, q_pos_emb)
+ key = apply_rotary_pos_emb(key, k_pos_emb)
+ else:
+ query_list = []
+ key_list = []
+ for i, rotary_pos_emb in enumerate(rotary_pos_emb_list):
+ rotary_pos_emb = [i[:, -cur_len:, :, :] for i in rotary_pos_emb]
+ rotary_pos_emb = (rotary_pos_emb,) * 2
+ q_pos_emb, k_pos_emb = rotary_pos_emb
+ # Slice the pos emb for current inference
+ query_list += [apply_rotary_pos_emb(query[i:i+1, :, :], q_pos_emb)]
+ key_list += [apply_rotary_pos_emb(key[i:i+1, :, :], k_pos_emb)]
+ query = torch.cat(query_list, dim=0)
+ key = torch.cat(key_list, dim=0)
+
+ if self.use_cache_quantization:
+ key = quantize_cache_v(key.permute(0, 2, 1, 3),
+ bits=8,
+ qmin=self.cache_qmin,
+ qmax=self.cache_qmax)
+ value = quantize_cache_v(value.permute(0, 2, 1, 3),
+ bits=8,
+ qmin=self.cache_qmin,
+ qmax=self.cache_qmax)
- if rotary_pos_emb is not None:
- q_pos_emb, k_pos_emb = rotary_pos_emb
- # Slice the pos emb for current inference
- cur_len = query.shape[1]
- q_pos_emb = q_pos_emb[:, -cur_len:, :, :]
- k_pos_emb = k_pos_emb[:, -cur_len:, :, :]
- query = apply_rotary_pos_emb(query, q_pos_emb)
- key = apply_rotary_pos_emb(key, k_pos_emb)
if layer_past is not None:
past_key, past_value = layer_past[0], layer_past[1]
- key = torch.cat((past_key, key), dim=1)
- value = torch.cat((past_value, value), dim=1)
+ if self.use_cache_quantization:
+ # use_cache_quantization:
+ # present=((q_key,key_scale,key_zero_point),
+ # (q_value,value_scale,value_zero_point))
+ key = (torch.cat((past_key[0], key[0]), dim=2),
+ torch.cat((past_key[1], key[1]), dim=2),
+ torch.cat((past_key[2], key[2]), dim=2))
+ value = (torch.cat((past_value[0], value[0]), dim=2),
+ torch.cat((past_value[1], value[1]), dim=2),
+ torch.cat((past_value[2], value[2]), dim=2))
+ else:
+ # not use_cache_quantization:
+ # present=(key,value)
+ key = torch.cat((past_key, key), dim=1)
+ value = torch.cat((past_value, value), dim=1)
if use_cache:
present = (key, value)
else:
present = None
- if self.use_logn_attn and not self.training:
- if self.logn_tensor.device != query.device or self.logn_tensor.dtype != query.dtype:
- self.logn_tensor = self.logn_tensor.to(query.device).type_as(query)
- seq_start = key.size(1) - query.size(1)
- seq_end = key.size(1)
- logn_tensor = self.logn_tensor[:, seq_start:seq_end, :, :]
+ key_size = key[0].size(2) if self.use_cache_quantization else key.size(1)
+ if key_size > self.seq_length and self.use_logn_attn and not self.training:
+ if self.use_cache_quantization:
+ seq_start = key[0].size(2) - query.size(1)
+ seq_end = key[0].size(2)
+ else:
+ seq_start = key.size(1) - query.size(1)
+ seq_end = key.size(1)
+ logn_tensor = self.logn_tensor[:, seq_start:seq_end, :, :].type_as(query)
query = query * logn_tensor.expand_as(query)
if (
@@ -434,23 +496,49 @@ def forward(
and query.is_cuda
):
q, k, v = query, key, value
- context_layer = self.core_attention_flash(q, k, v)
-
- context_layer = rearrange(
- context_layer, "b s h d -> b s (h d)"
- ).contiguous()
+ attn_output = self.core_attention_flash(q, k, v, attention_mask=attention_mask)
else:
+ key_size = key[0].size(2) if self.use_cache_quantization else key.size(1)
+ if query.size(1) == key_size:
+ causal_mask = torch.tril(
+ torch.ones((key_size, key_size), dtype=torch.bool, device=query.device)
+ ).view(1, 1, key_size, key_size)
+ else:
+ causal_mask = None
query = query.permute(0, 2, 1, 3)
- key = key.permute(0, 2, 1, 3)
- value = value.permute(0, 2, 1, 3)
- attn_output, attn_weight = self._attn(
- query, key, value, attention_mask, head_mask
- )
- context_layer = self._merge_heads(
- attn_output, self.num_heads, self.head_dim
- )
+ if not self.use_cache_quantization:
+ key = key.permute(0, 2, 1, 3)
+ value = value.permute(0, 2, 1, 3)
+ if (
+ causal_mask is None
+ and self.use_flash_attn
+ and flash_attn_unpadded_func is not None
+ and not self.is_fp32
+ and not query.is_cuda
+ ):
+ raise Exception(_ERROR_INPUT_CPU_QUERY_WITH_FLASH_ATTN_ACTIVATED)
+
+ if not self.use_cache_quantization and SUPPORT_TORCH2:
+ if attention_mask is not None:
+ attention_mask = attention_mask.expand(-1, -1, query.size(2), -1)
+ if causal_mask is not None:
+ attention_mask = attention_mask.masked_fill(~causal_mask, torch.finfo(query.dtype).min)
+ else:
+ attention_mask = causal_mask
+ attn_output = F.scaled_dot_product_attention(
+ query, key, value, attn_mask=attention_mask
+ ).transpose(1, 2)
+ attn_weight = None
+ else:
+ attn_output, attn_weight = self._attn(
+ query, key, value, causal_mask, attention_mask, head_mask
+ )
+ context_layer = self._merge_heads(
+ attn_output, self.num_heads, self.head_dim
+ )
attn_output = self.c_proj(context_layer)
+
outputs = (attn_output, present)
if output_attentions:
if (
@@ -459,6 +547,8 @@ def forward(
and not self.is_fp32
):
raise ValueError("Cannot output attentions while using flash-attn")
+ elif not self.use_cache_quantization and SUPPORT_TORCH2:
+ raise ValueError("Cannot output attentions while using scaled_dot_product_attention")
else:
outputs += (attn_weight,)
@@ -469,12 +559,12 @@ class QWenMLP(nn.Module):
def __init__(self, config):
super().__init__()
self.w1 = nn.Linear(
- config.hidden_size, config.ffn_hidden_size // 2, bias=not config.no_bias
+ config.hidden_size, config.intermediate_size // 2, bias=not config.no_bias
)
self.w2 = nn.Linear(
- config.hidden_size, config.ffn_hidden_size // 2, bias=not config.no_bias
+ config.hidden_size, config.intermediate_size // 2, bias=not config.no_bias
)
- ff_dim_in = config.ffn_hidden_size // 2
+ ff_dim_in = config.intermediate_size // 2
self.c_proj = nn.Linear(ff_dim_in, config.hidden_size, bias=not config.no_bias)
def forward(self, hidden_states):
@@ -486,24 +576,16 @@ def forward(self, hidden_states):
class QWenBlock(nn.Module):
- def __init__(self, config, layer_idx=None, num_expert=1):
+ def __init__(self, config):
super().__init__()
- self.num_expert = num_expert
- self.layer_number = layer_idx
- self.apply_residual_connection_post_layernorm = (
- config.apply_residual_connection_post_layernorm
- )
hidden_size = config.hidden_size
- self.apply_residual_connection_post_layernorm = (
- config.apply_residual_connection_post_layernorm
- )
self.bf16 = config.bf16
self.ln_1 = RMSNorm(
hidden_size,
eps=config.layer_norm_epsilon,
)
- self.attn = QWenAttention(config, layer_number=layer_idx)
+ self.attn = QWenAttention(config)
self.ln_2 = RMSNorm(
hidden_size,
eps=config.layer_norm_epsilon,
@@ -514,6 +596,7 @@ def __init__(self, config, layer_idx=None, num_expert=1):
def forward(
self,
hidden_states: Optional[Tuple[torch.FloatTensor]],
+ rotary_pos_emb_list: Optional[List[List[torch.Tensor]]] = None,
layer_past: Optional[Tuple[torch.Tensor]] = None,
attention_mask: Optional[torch.FloatTensor] = None,
head_mask: Optional[torch.FloatTensor] = None,
@@ -526,6 +609,7 @@ def forward(
attn_outputs = self.attn(
layernorm_output,
+ rotary_pos_emb_list,
layer_past=layer_past,
attention_mask=attention_mask,
head_mask=head_mask,
@@ -536,19 +620,12 @@ def forward(
outputs = attn_outputs[1:]
- if self.apply_residual_connection_post_layernorm:
- residual = layernorm_output
- else:
- residual = hidden_states
+ residual = hidden_states
layernorm_input = attn_output + residual
layernorm_output = self.ln_2(layernorm_input)
- if self.apply_residual_connection_post_layernorm:
- residual = layernorm_output
- else:
- residual = layernorm_input
-
+ residual = layernorm_input
mlp_output = self.mlp(layernorm_output)
hidden_states = residual + mlp_output
@@ -566,6 +643,7 @@ class QWenPreTrainedModel(PreTrainedModel):
is_parallelizable = False
supports_gradient_checkpointing = True
_no_split_modules = ["QWenBlock"]
+ _skip_keys_device_placement = "past_key_values"
def __init__(self, *inputs, **kwargs):
super().__init__(*inputs, **kwargs)
@@ -589,7 +667,7 @@ def _init_weights(self, module):
mean=0.0,
std=(
self.config.initializer_range
- / math.sqrt(2 * self.config.n_layer)
+ / math.sqrt(2 * self.config.num_hidden_layers)
),
)
@@ -603,31 +681,40 @@ class QWenModel(QWenPreTrainedModel):
def __init__(self, config):
super().__init__(config)
- self.vocab_size = config.padded_vocab_size
+ self.vocab_size = config.vocab_size
self.num_hidden_layers = config.num_hidden_layers
self.embed_dim = config.hidden_size
+ self.use_cache_quantization = self.config.use_cache_quantization if hasattr(self.config, 'use_cache_quantization') else False
- max_sequence_length = config.max_position_embeddings
- self.position_embedding_type = config.pos_emb
self.gradient_checkpointing = False
+ self.use_dynamic_ntk = config.use_dynamic_ntk
+ self.seq_length = config.seq_length
+
+ self.wte = nn.Embedding(self.vocab_size, self.embed_dim)
- if self.position_embedding_type == "learned":
- self.wpe = nn.Embedding(max_sequence_length, self.embed_dim)
- self.init_method(self.position_embeddings.weight)
- self._position_embeddings_key = "position_embeddings"
- self.init_method(self.position_embeddings.weight)
+ self.drop = nn.Dropout(config.emb_dropout_prob)
+
+ if config.rotary_pct == 1.0:
+ self.rotary_ndims = None
else:
- self.wpe = None
- self._position_embeddings_key = ""
+ assert config.rotary_pct < 1
+ self.rotary_ndims = int(
+ config.kv_channels * config.rotary_pct
+ )
+ dim = (
+ self.rotary_ndims
+ if self.rotary_ndims is not None
+ else config.kv_channels
+ )
+ self.rotary_emb = RotaryEmbedding(dim, base=config.rotary_emb_base)
- self.wte = nn.Embedding(self.vocab_size, self.embed_dim)
+ self.use_flash_attn = config.use_flash_attn
+ self.is_fp32 = not (config.bf16 or config.fp16)
- self.drop = nn.Dropout(config.embd_pdrop)
self.h = nn.ModuleList(
[
QWenBlock(
- config,
- layer_idx=i,
+ config
)
for i in range(config.num_hidden_layers)
]
@@ -645,6 +732,12 @@ def get_input_embeddings(self):
def set_input_embeddings(self, new_embeddings):
self.wte = new_embeddings
+ def get_ntk_alpha(self, true_seq_len):
+ context_value = math.log(true_seq_len / self.seq_length, 2) + 1
+ ntk_alpha = 2 ** math.ceil(context_value) - 1
+ ntk_alpha = max(ntk_alpha, 1)
+ return ntk_alpha
+
def forward(
self,
input_ids: Optional[torch.LongTensor] = None,
@@ -701,8 +794,10 @@ def forward(
past_length = 0
past_key_values = tuple([None] * len(self.h))
else:
- past_length = past_key_values[0][0].size(-2)
-
+ if self.use_cache_quantization:
+ past_length = past_key_values[0][0][0].size(2)
+ else:
+ past_length = past_key_values[0][0].size(-2)
if position_ids is None:
position_ids = torch.arange(
past_length,
@@ -719,17 +814,41 @@ def forward(
attention_mask = attention_mask[:, None, None, :]
attention_mask = attention_mask.to(dtype=self.dtype)
attention_mask = (1.0 - attention_mask) * torch.finfo(self.dtype).min
- # attention_mask中mask掉的部分是-inf, 看到的部分是0
encoder_attention_mask = None
- head_mask = self.get_head_mask(head_mask, self.config.n_layer)
+ head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
if inputs_embeds is None:
inputs_embeds = self.wte(input_ids)
hidden_states = inputs_embeds
- if self.wpe is not None:
- position_embeds = self.wpe(position_ids)
- hidden_states = hidden_states + position_embeds
+
+ kv_seq_len = hidden_states.size()[1]
+ if past_key_values[0] is not None:
+ # past key values[0][0] shape: bs * seq_len * head_num * dim
+ if self.use_cache_quantization:
+ kv_seq_len += past_key_values[0][0][0].shape[2]
+ else:
+ kv_seq_len += past_key_values[0][0].shape[1]
+
+ if self.training or not self.use_dynamic_ntk:
+ ntk_alpha_list = [1.0]
+ elif kv_seq_len != hidden_states.size()[1]:
+ ntk_alpha_list = self.rotary_emb._ntk_alpha_cached_list
+ else:
+ ntk_alpha_list = []
+ if attention_mask is not None and kv_seq_len > self.seq_length:
+ true_seq_lens = attention_mask.squeeze(1).squeeze(1).eq(0).sum(dim=-1, dtype=torch.int32)
+ for i in range(hidden_states.size()[0]):
+ true_seq_len = true_seq_lens[i].item()
+ ntk_alpha = self.get_ntk_alpha(true_seq_len)
+ ntk_alpha_list.append(ntk_alpha)
+ else:
+ ntk_alpha = self.get_ntk_alpha(kv_seq_len)
+ ntk_alpha_list.append(ntk_alpha)
+ self.rotary_emb._ntk_alpha_cached_list = ntk_alpha_list
+ rotary_pos_emb_list = [
+ self.rotary_emb(kv_seq_len, ntk_alpha=ntk_alpha) for ntk_alpha in ntk_alpha_list
+ ]
hidden_states = self.drop(hidden_states)
output_shape = input_shape + (hidden_states.size(-1),)
@@ -761,6 +880,7 @@ def custom_forward(*inputs):
outputs = torch.utils.checkpoint.checkpoint(
create_custom_forward(block),
hidden_states,
+ rotary_pos_emb_list,
None,
attention_mask,
head_mask[i],
@@ -771,6 +891,7 @@ def custom_forward(*inputs):
outputs = block(
hidden_states,
layer_past=layer_past,
+ rotary_pos_emb_list=rotary_pos_emb_list,
attention_mask=attention_mask,
head_mask=head_mask[i],
encoder_hidden_states=encoder_hidden_states,
@@ -781,13 +902,16 @@ def custom_forward(*inputs):
hidden_states = outputs[0]
if use_cache is True:
- presents = presents + (outputs[2 if output_attentions else 1],)
+ presents = presents + (outputs[1],)
if output_attentions:
- all_self_attentions = all_self_attentions + (outputs[1],)
+ all_self_attentions = all_self_attentions + (outputs[2 if use_cache else 1],)
hidden_states = self.ln_f(hidden_states)
hidden_states = hidden_states.view(output_shape)
+ # Add last hidden state
+ if output_hidden_states:
+ all_hidden_states = all_hidden_states + (hidden_states,)
if not return_dict:
return tuple(
@@ -839,7 +963,7 @@ def __init__(self, config):
logger.warn("Your device support faster inference by passing bf16=True in \"AutoModelForCausalLM.from_pretrained\".")
elif SUPPORT_FP16:
logger.warn("Your device support faster inference by passing fp16=True in \"AutoModelForCausalLM.from_pretrained\".")
-
+
if config.use_flash_attn == "auto":
if config.bf16 or config.fp16:
logger.warn("Try importing flash-attention for faster inference...")
@@ -853,7 +977,7 @@ def __init__(self, config):
_import_flash_attn()
self.transformer = QWenModel(config)
- self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
if config.bf16:
self.transformer.bfloat16()
@@ -872,22 +996,13 @@ def set_output_embeddings(self, new_embeddings):
def prepare_inputs_for_generation(
self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs
):
- token_type_ids = kwargs.get("token_type_ids", None)
if past_key_values:
input_ids = input_ids[:, -1].unsqueeze(-1)
- if token_type_ids is not None:
- token_type_ids = token_type_ids[:, -1].unsqueeze(-1)
-
- attention_mask = kwargs.get("attention_mask", None)
- position_ids = kwargs.get("position_ids", None)
- if attention_mask is not None and position_ids is None:
- position_ids = attention_mask.long().cumsum(-1) - 1
- position_ids.masked_fill_(attention_mask == 0, 1)
- if past_key_values:
- position_ids = position_ids[:, -1].unsqueeze(-1)
+ if input_ids.size(0) == 1:
+ attention_mask = None
else:
- position_ids = None
+ attention_mask = kwargs.get("attention_mask", None)
if inputs_embeds is not None and past_key_values is None:
model_inputs = {"inputs_embeds": inputs_embeds}
@@ -898,9 +1013,7 @@ def prepare_inputs_for_generation(
{
"past_key_values": past_key_values,
"use_cache": kwargs.get("use_cache"),
- "position_ids": position_ids,
"attention_mask": attention_mask,
- "token_type_ids": token_type_ids,
}
)
return model_inputs
@@ -987,35 +1100,45 @@ def chat(
query: str,
history: Optional[HistoryType],
system: str = "You are a helpful assistant.",
- append_history: bool = True,
stream: Optional[bool] = _SENTINEL,
stop_words_ids: Optional[List[List[int]]] = None,
+ generation_config: Optional[GenerationConfig] = None,
**kwargs,
) -> Tuple[str, HistoryType]:
+ generation_config = generation_config if generation_config is not None else self.generation_config
+
assert stream is _SENTINEL, _ERROR_STREAM_IN_CHAT
- assert self.generation_config.chat_format == 'chatml', _ERROR_BAD_CHAT_FORMAT
+ assert generation_config.chat_format == 'chatml', _ERROR_BAD_CHAT_FORMAT
if history is None:
history = []
+ else:
+ # make a copy of the user's input such that is is left untouched
+ history = copy.deepcopy(history)
+
if stop_words_ids is None:
stop_words_ids = []
+ max_window_size = kwargs.get('max_window_size', None)
+ if max_window_size is None:
+ max_window_size = generation_config.max_window_size
raw_text, context_tokens = make_context(
tokenizer,
query,
history=history,
system=system,
- max_window_size=6144,
- chat_format=self.generation_config.chat_format,
+ max_window_size=max_window_size,
+ chat_format=generation_config.chat_format,
)
stop_words_ids.extend(get_stop_words_ids(
- self.generation_config.chat_format, tokenizer
+ generation_config.chat_format, tokenizer
))
input_ids = torch.tensor([context_tokens]).to(self.device)
outputs = self.generate(
input_ids,
- stop_words_ids = stop_words_ids,
- return_dict_in_generate = False,
+ stop_words_ids=stop_words_ids,
+ return_dict_in_generate=False,
+ generation_config=generation_config,
**kwargs,
)
@@ -1024,13 +1147,16 @@ def chat(
tokenizer,
raw_text_len=len(raw_text),
context_length=len(context_tokens),
- chat_format=self.generation_config.chat_format,
+ chat_format=generation_config.chat_format,
verbose=False,
errors='replace'
)
- if append_history:
- history.append((query, response))
+ # as history is a copy of the user inputs,
+ # we can always return the new turn to the user.
+ # separating input history and output history also enables the user
+ # to implement more complex history management
+ history.append((query, response))
return response, history
@@ -1042,30 +1168,35 @@ def chat_stream(
system: str = "You are a helpful assistant.",
stop_words_ids: Optional[List[List[int]]] = None,
logits_processor: Optional[LogitsProcessorList] = None,
+ generation_config: Optional[GenerationConfig] = None,
**kwargs,
) -> Generator[str, Any, None]:
- assert self.generation_config.chat_format == 'chatml', _ERROR_BAD_CHAT_FORMAT
+ generation_config = generation_config if generation_config is not None else self.generation_config
+ assert generation_config.chat_format == 'chatml', _ERROR_BAD_CHAT_FORMAT
if history is None:
history = []
if stop_words_ids is None:
stop_words_ids = []
+ max_window_size = kwargs.get('max_window_size', None)
+ if max_window_size is None:
+ max_window_size = generation_config.max_window_size
raw_text, context_tokens = make_context(
tokenizer,
query,
history=history,
system=system,
- max_window_size=6144,
- chat_format=self.generation_config.chat_format,
+ max_window_size=max_window_size,
+ chat_format=generation_config.chat_format,
)
stop_words_ids.extend(get_stop_words_ids(
- self.generation_config.chat_format, tokenizer
+ generation_config.chat_format, tokenizer
))
if stop_words_ids is not None:
stop_words_logits_processor = StopWordsLogitsProcessor(
stop_words_ids=stop_words_ids,
- eos_token_id=self.generation_config.eos_token_id,
+ eos_token_id=generation_config.eos_token_id,
)
if logits_processor is None:
logits_processor = LogitsProcessorList([stop_words_logits_processor])
@@ -1076,7 +1207,8 @@ def chat_stream(
from transformers_stream_generator.main import NewGenerationMixin, StreamGenerationConfig
self.__class__.generate_stream = NewGenerationMixin.generate
self.__class__.sample_stream = NewGenerationMixin.sample_stream
- stream_config = StreamGenerationConfig(**self.generation_config.to_dict(), do_stream=True)
+ stream_config = StreamGenerationConfig(**generation_config.to_dict(), do_stream=True)
+
def stream_generator():
outputs = []
for token in self.generate_stream(
@@ -1105,17 +1237,19 @@ def generate(
streamer: Optional["BaseStreamer"] = None,
**kwargs,
) -> Union[GenerateOutput, torch.LongTensor]:
+ generation_config = generation_config if generation_config is not None else self.generation_config
+
# Process stop_words_ids.
stop_words_ids = kwargs.pop("stop_words_ids", None)
if stop_words_ids is None and generation_config is not None:
stop_words_ids = getattr(generation_config, "stop_words_ids", None)
if stop_words_ids is None:
- stop_words_ids = getattr(self.generation_config, "stop_words_ids", None)
+ stop_words_ids = getattr(generation_config, "stop_words_ids", None)
if stop_words_ids is not None:
stop_words_logits_processor = StopWordsLogitsProcessor(
stop_words_ids=stop_words_ids,
- eos_token_id=self.generation_config.eos_token_id,
+ eos_token_id=generation_config.eos_token_id,
)
if logits_processor is None:
logits_processor = LogitsProcessorList([stop_words_logits_processor])
@@ -1140,16 +1274,17 @@ def __init__(self, dim, base=10000):
super().__init__()
self.dim = dim
self.base = base
- self.inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+ inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
if importlib.util.find_spec("einops") is None:
raise RuntimeError("einops is required for Rotary Embedding")
self._rotary_pos_emb_cache = None
self._seq_len_cached = 0
self._ntk_alpha_cached = 1.0
+ self._ntk_alpha_cached_list = [1.0]
- def update_rotary_pos_emb_cache(self, max_seq_len, offset=0, ntk_alpha=1.0):
- seqlen = max_seq_len + offset
+ def update_rotary_pos_emb_cache(self, seqlen, ntk_alpha=1.0):
if seqlen > self._seq_len_cached or ntk_alpha != self._ntk_alpha_cached:
base = self.base * ntk_alpha ** (self.dim / (self.dim - 2))
self.inv_freq = 1.0 / (
@@ -1163,14 +1298,19 @@ def update_rotary_pos_emb_cache(self, max_seq_len, offset=0, ntk_alpha=1.0):
self._ntk_alpha_cached = ntk_alpha
seq = torch.arange(self._seq_len_cached, device=self.inv_freq.device)
freqs = torch.outer(seq.type_as(self.inv_freq), self.inv_freq)
+
emb = torch.cat((freqs, freqs), dim=-1)
from einops import rearrange
- self._rotary_pos_emb_cache = rearrange(emb, "n d -> 1 n 1 d")
+ emb = rearrange(emb, "n d -> 1 n 1 d")
- def forward(self, max_seq_len, offset=0, ntk_alpha=1.0):
- self.update_rotary_pos_emb_cache(max_seq_len, offset, ntk_alpha)
- return self._rotary_pos_emb_cache[:, offset : offset + max_seq_len]
+ cos, sin = emb.cos(), emb.sin()
+ self._rotary_pos_emb_cache = [cos, sin]
+
+ def forward(self, max_seq_len, ntk_alpha=1.0):
+ self.update_rotary_pos_emb_cache(max_seq_len, ntk_alpha)
+ cos, sin = self._rotary_pos_emb_cache
+ return [cos[:, :max_seq_len], sin[:, :max_seq_len]]
def _rotate_half(x):
@@ -1182,20 +1322,28 @@ def _rotate_half(x):
def apply_rotary_pos_emb(t, freqs):
- if apply_rotary_emb_func is not None:
- t_ = t.float()
- freqs = freqs.squeeze(0).squeeze(1)
- cos = freqs[:, : freqs.shape[-1] // 2].cos()
- sin = freqs[:, : freqs.shape[-1] // 2].sin()
- output = apply_rotary_emb_func(t_, cos, sin).type_as(t)
- return output
+ """ Apply rotary embedding to the first rotary_dim of the iput
+
+ Arguments:
+ t (tensor(batch_size, seq_len, n_head, head_dim)):
+ the input embedding/hidden states
+ freqs (list[tensor(1, seq_len, 1, rotary_dim), tensor(1, seq_len, 1, rotary_dim)]):
+ the cached cos/sin position embeddings
+ """
+ rot_dim = freqs[0].shape[-1]
+ cos, sin = freqs
+ t_float = t.float()
+ if apply_rotary_emb_func is not None and t.is_cuda:
+ # apply_rotary_emb in flash_attn requires cos/sin to be of
+ # shape (seqlen, rotary_dim / 2) and apply rotary embedding
+ # to the first rotary_dim of the input
+ cos = cos.squeeze(0).squeeze(1)[:, : rot_dim // 2]
+ sin = sin.squeeze(0).squeeze(1)[:, : rot_dim // 2]
+ return apply_rotary_emb_func(t_float, cos, sin).type_as(t)
else:
- rot_dim = freqs.shape[-1]
- t_, t_pass_ = t[..., :rot_dim], t[..., rot_dim:]
- t_ = t_.float()
- t_pass_ = t_pass_.float()
- t_ = (t_ * freqs.cos()) + (_rotate_half(t_) * freqs.sin())
- return torch.cat((t_, t_pass_), dim=-1).type_as(t)
+ t_rot, t_pass = t_float[..., :rot_dim], t_float[..., rot_dim:]
+ t_rot = (t_rot * cos) + (_rotate_half(t_rot) * sin)
+ return torch.cat((t_rot, t_pass), dim=-1).type_as(t)
class RMSNorm(torch.nn.Module):
diff --git a/mft_peft_hf/src/model/qwen/qwen_generation_utils.py b/mftcoder_accelerate/src/model/qwen/qwen_generation_utils.py
similarity index 100%
rename from mft_peft_hf/src/model/qwen/qwen_generation_utils.py
rename to mftcoder_accelerate/src/model/qwen/qwen_generation_utils.py
diff --git a/mft_peft_hf/src/model/qwen/tokenization_qwen.py b/mftcoder_accelerate/src/model/qwen/tokenization_qwen.py
similarity index 76%
rename from mft_peft_hf/src/model/qwen/tokenization_qwen.py
rename to mftcoder_accelerate/src/model/qwen/tokenization_qwen.py
index 4a66a09..2a526d6 100644
--- a/mft_peft_hf/src/model/qwen/tokenization_qwen.py
+++ b/mftcoder_accelerate/src/model/qwen/tokenization_qwen.py
@@ -27,11 +27,22 @@
# regular texts, the surface forms of special tokens need to be
# as different as possible to minimize the impact
EXTRAS = tuple((f"<|extra_{i}|>" for i in range(205)))
-SPECIAL_TOKENS = (
- ENDOFTEXT,
- IMSTART,
- IMEND,
-) + EXTRAS
+# changed to use actual index to avoid misconfiguration with vocabulary expansion
+SPECIAL_START_ID = 151643
+SPECIAL_TOKENS = tuple(
+ enumerate(
+ (
+ (
+ ENDOFTEXT,
+ IMSTART,
+ IMEND,
+ )
+ + EXTRAS
+ ),
+ start=SPECIAL_START_ID,
+ )
+)
+SPECIAL_TOKENS_SET = set(t for i, t in SPECIAL_TOKENS)
def _load_tiktoken_bpe(tiktoken_bpe_file: str) -> Dict[bytes, int]:
@@ -42,6 +53,7 @@ def _load_tiktoken_bpe(tiktoken_bpe_file: str) -> Dict[bytes, int]:
for token, rank in (line.split() for line in contents.splitlines() if line)
}
+
class QWenTokenizer(PreTrainedTokenizer):
"""QWen tokenizer."""
@@ -51,20 +63,35 @@ def __init__(
self,
vocab_file,
errors="replace",
+ extra_vocab_file=None,
**kwargs,
):
super().__init__(**kwargs)
- self.errors = errors # how to handle errors in decoding
+ # how to handle errors in decoding UTF-8 byte sequences
+ # use ignore if you are in streaming inference
+ self.errors = errors
- self.mergeable_ranks = _load_tiktoken_bpe(vocab_file) # type: dict[bytes, int]
+ self.mergeable_ranks = _load_tiktoken_bpe(vocab_file) # type: Dict[bytes, int]
self.special_tokens = {
token: index
- for index, token in enumerate(
- SPECIAL_TOKENS, start=len(self.mergeable_ranks)
- )
+ for index, token in SPECIAL_TOKENS
}
+ # try load extra vocab from file
+ if extra_vocab_file is not None:
+ used_ids = set(self.mergeable_ranks.values()) | set(self.special_tokens.values())
+ extra_mergeable_ranks = _load_tiktoken_bpe(extra_vocab_file)
+ for token, index in extra_mergeable_ranks.items():
+ if token in self.mergeable_ranks:
+ logger.info(f"extra token {token} exists, skipping")
+ continue
+ if index in used_ids:
+ logger.info(f'the index {index} for extra token {token} exists, skipping')
+ continue
+ self.mergeable_ranks[token] = index
+ # the index may be sparse after this, but don't worry tiktoken.Encoding will handle this
+
enc = tiktoken.Encoding(
"Qwen",
pat_str=PAT_STR,
@@ -86,6 +113,23 @@ def __init__(
self.im_start_id = self.special_tokens[IMSTART]
self.im_end_id = self.special_tokens[IMEND]
+ def __getstate__(self):
+ # for pickle lovers
+ state = self.__dict__.copy()
+ del state["tokenizer"]
+ return state
+
+ def __setstate__(self, state):
+ # tokenizer is not python native; don't pass it; rebuild it
+ self.__dict__.update(state)
+ enc = tiktoken.Encoding(
+ "Qwen",
+ pat_str=PAT_STR,
+ mergeable_ranks=self.mergeable_ranks,
+ special_tokens=self.special_tokens,
+ )
+ self.tokenizer = enc
+
def __len__(self) -> int:
return self.tokenizer.n_vocab
@@ -108,13 +152,17 @@ def convert_tokens_to_ids(
ids.append(self.mergeable_ranks.get(token))
return ids
- def _add_tokens(self, new_tokens: Union[List[str], List[AddedToken]], special_tokens: bool = False) -> int:
+ def _add_tokens(
+ self,
+ new_tokens: Union[List[str], List[AddedToken]],
+ special_tokens: bool = False,
+ ) -> int:
if not special_tokens and new_tokens:
- raise ValueError('Adding regular tokens is not supported')
+ raise ValueError("Adding regular tokens is not supported")
for token in new_tokens:
surface_form = token.content if isinstance(token, AddedToken) else token
- if surface_form not in SPECIAL_TOKENS:
- raise ValueError('Adding unknown special tokens is not supported')
+ if surface_form not in SPECIAL_TOKENS_SET:
+ raise ValueError("Adding unknown special tokens is not supported")
return 0
def save_vocabulary(self, save_directory: str, **kwargs) -> Tuple[str]:
diff --git a/mftcoder_accelerate/src/model/qwen/tokenizer_config.json b/mftcoder_accelerate/src/model/qwen/tokenizer_config.json
new file mode 100644
index 0000000..9c37cac
--- /dev/null
+++ b/mftcoder_accelerate/src/model/qwen/tokenizer_config.json
@@ -0,0 +1,10 @@
+{
+ "model_max_length": 8192,
+ "tokenizer_class": "QWenTokenizer",
+ "auto_map": {
+ "AutoTokenizer": [
+ "tokenization_qwen.QWenTokenizer",
+ null
+ ]
+ }
+}
diff --git a/mftcoder_accelerate/src/mpt/mpt_accelerate.py b/mftcoder_accelerate/src/mpt/mpt_accelerate.py
new file mode 100644
index 0000000..5d187c9
--- /dev/null
+++ b/mftcoder_accelerate/src/mpt/mpt_accelerate.py
@@ -0,0 +1,494 @@
+"""
+# @author Chaoyu Chen
+# @date 2024/6/1
+# @module mpt_accelerate.py
+
+Accelerate + DeepSpeed + Full-parameter + Multi-task + Pre-training/Continue Training/Finetuning
+
+Entry
+"""
+
+import os
+import sys
+import argparse
+import math
+import logging
+import json
+import time
+from tqdm.auto import tqdm
+import transformers
+import numpy as np
+import torch
+from torch import nn
+from dataclasses import dataclass
+from datasets import Dataset
+import datasets
+from torch.utils.data import DataLoader
+from torch.distributed.fsdp.fully_sharded_data_parallel import FullOptimStateDictConfig, FullStateDictConfig
+
+from transformers import (
+ AutoModelForCausalLM,
+ AutoTokenizer,
+ get_linear_schedule_with_warmup,
+ set_seed,
+ BitsAndBytesConfig,
+ get_scheduler,
+)
+
+from accelerate import Accelerator, DistributedType, FullyShardedDataParallelPlugin, DataLoaderConfiguration
+from accelerate.logging import get_logger
+from datetime import timedelta
+from accelerate.utils import InitProcessGroupKwargs
+from transformers.optimization import Adafactor
+
+# insert src as import path
+current_path = os.path.abspath(__file__)
+parent_dir = os.path.dirname(os.path.dirname(current_path))
+sys.path.insert(0, parent_dir)
+
+from tokenizer import build_tokenizer
+from data.multi_task_dataset import load_dataset_from_jsonl, compile_helper
+from data.data_utils import load_dataset_from_bin
+from utils.common_utils import print_rank_0, generate_task_id, TASK2ID, ID2TASK
+from mpt.mpt_trainer import MptTrainer
+from mpt.mpt_arguments import MptTrainArgs
+from utils.model_mapping import MODEL_TYPES, SUPPORT_IN_TRANSFORMERS
+
+
+logger = get_logger(__name__)
+
+
+def get_task_mask(args, task_id):
+ task_num = len(TASK2ID)
+ task_mask = torch.zeros(task_id.shape[0], task_num)
+ task_mask[torch.arange(task_id.size(0)).unsqueeze(1), task_id] = 1
+
+ return task_mask
+
+
+def get_attention_mask_and_position_ids(data):
+ """Build masks and position id for left to right model."""
+
+ # Extract batch size and sequence length.
+ batch_size, seq_length = data.size()
+
+ attention_mask = torch.ones((batch_size, seq_length), device=data.device)
+
+ # Position ids.
+ position_ids = torch.arange(seq_length, dtype=torch.long, device=data.device)
+ position_ids = position_ids.unsqueeze(0).expand_as(data).clone()
+
+ return attention_mask, position_ids
+
+
+@dataclass
+class DataCollatorForMFTDataset(object):
+ args: None
+
+ def __call__(self, instances):
+ (input_ids, loss_mask, weights, task_id) = tuple(
+ [instance.get(key, None) for instance in instances]
+ for key in ("input_ids", "loss_mask", "weight", "task_id")
+ )
+
+ result_batch = {}
+ """
+ outputs = model(
+ input_ids=batch['input_ids'],
+ attention_mask=batch['attention_mask'],
+ # labels=(batch['labels'], batch['loss_mask'], batch['task_mask']),
+ # labels=(batch['labels'], batch['loss_mask']),
+ position_ids=batch['position_ids'])
+ """
+
+ # if loss_mask is not None:
+ loss_mask = torch.tensor(np.array(loss_mask)).long()
+ last_one_pos = (loss_mask == 1).long().cumsum(dim=1).argmax(dim=1)
+ if self.args.use_dynamic_padding:
+ # get last non-padding position
+ max_pos = last_one_pos.max().item() + 1
+ else:
+ max_pos = loss_mask.shape[-1]
+
+ if self.args.tokenize_mode == "sst" and self.args.padding_mode == "pack":
+ # 兼容sst + pack tokenization, 最后一位是脏数据，需要去掉
+ result_batch["loss_mask"] = loss_mask.float()[:, 1 : max_pos - 1].contiguous()
+ input_ids = torch.tensor(np.array(input_ids)).long()
+ result_batch["input_ids"] = input_ids[:, : max_pos - 2].contiguous()
+ result_batch["labels"] = input_ids[:, 1 : max_pos - 1].contiguous()
+ else:
+ result_batch["loss_mask"] = loss_mask.float()[:, 1:max_pos].contiguous()
+ input_ids = torch.tensor(np.array(input_ids)).long()
+ # print(f"shape of input_ids: {input_ids.shape}")
+ result_batch["input_ids"] = input_ids[:, : max_pos - 1].contiguous()
+ result_batch["labels"] = input_ids[:, 1:max_pos].contiguous()
+
+ # Get the masks and position ids.
+
+ # if you want to be compatible with non-gpt models, something you can do here
+ if self.args.model_type in ["antglm"]:
+ (result_batch["attention_mask"], result_batch["position_ids"]) = get_attention_mask_and_position_ids(
+ data=result_batch["input_ids"]
+ )
+ elif self.args.model_type in ["mixtral", "mtx-qwen2", "qwen2_moe"]:
+ batch_size, seq_length = result_batch["input_ids"].shape
+ # bsz * seq_length
+ range_tensor = torch.arange(seq_length).unsqueeze(0).repeat(batch_size, 1)
+ # attention_mask for padding tokens
+ attention_mask = (range_tensor <= last_one_pos.reshape(batch_size, 1)).long()
+ result_batch["attention_mask"], result_batch["position_ids"] = attention_mask, None
+ else:
+ # For decoder-only models, transformers will create them.
+ result_batch["attention_mask"], result_batch["position_ids"] = None, None
+
+ if task_id is not None:
+ task_id = torch.tensor(np.array(task_id))
+ result_batch["task_mask"] = get_task_mask(self.args, task_id) # bsz * task_num
+ result_batch["task_id"] = task_id
+
+ return result_batch
+
+
+def pprint_args(args, accelerator):
+ # 计算所有键的最大字符串长度
+ max_key_length = max(len(str(key)) for key in vars(args).keys())
+
+ message = ""
+ message += "====" * 60 + "\n"
+ message += "\n".join([f"{k:<{max_key_length}} : {v}" for k, v in vars(args).items()]) + "\n"
+ message += "====" * 60 + "\n"
+ accelerator.print(message)
+ accelerator.print("GPU: {}".format(torch.cuda.current_device()))
+
+
+def prepare_args():
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--train_config", type=str, default=None)
+
+ parser.add_argument("--data_paths", type=str, default=None)
+ parser.add_argument("--output_dir", type=str, default=None)
+ parser.add_argument("--tb_dir", type=str, default=None)
+ parser.add_argument("--pretrained_model_path", type=str, default=None)
+ parser.add_argument("--micro_batch_size", type=int, default=None)
+ parser.add_argument("--model_type", type=str, default=None)
+ parser.add_argument("--distributed_type", type=str, default="deepspeed")
+
+ parsed = parser.parse_args()
+ # get json configs
+ with open(parsed.train_config, "r") as f:
+ train_config = json.load(f)
+
+ # parse args from cofig.json
+ # args = argparse.Namespace(**train_config)
+ args = MptTrainArgs(**train_config)
+
+ # override args by cli arguments
+ if parsed.data_paths:
+ args.data_paths = parsed.data_paths
+ if parsed.output_dir:
+ args.output_dir = parsed.output_dir
+ if parsed.tb_dir:
+ args.tb_dir = parsed.tb_dir
+ if parsed.pretrained_model_path:
+ args.pretrained_model_path = parsed.pretrained_model_path
+ args.vocab_file = parsed.pretrained_model_path
+ if parsed.micro_batch_size:
+ args.per_device_train_batch_size = parsed.micro_batch_size
+ args.per_device_eval_batch_size = parsed.micro_batch_size
+ if parsed.model_type:
+ args.model_type = parsed.model_type
+
+ args.distributed_type = parsed.distributed_type
+
+ # refactor args
+
+ args.vocab_file = args.pretrained_model_path
+
+ args.data_weights = "[" + ",".join(["1."] * len(args.data_paths[1:-1].split(","))) + "]"
+
+ # generate TASK2ID, ID2TASK
+ generate_task_id(args.data_paths)
+
+ if args.weighted_loss_mode == "coba":
+ args.task_weights = [1.0] * len(ID2TASK)
+ elif args.task_weights is not None:
+ args.task_weights = [float(wt) for wt in args.task_weights[1:-1].split(",")]
+ assert len(args.task_weights) == len(ID2TASK), f"length of task_weights must equal to length of data_paths"
+ else:
+ args.task_weights = [1.0] * len(ID2TASK)
+
+ return args
+
+
+def main():
+ t0 = time.time()
+ os.environ["TOKENIZERS_PARALLELISM"] = "false"
+ os.environ["HF_HUB_OFFLINE"] = "false"
+ # get input args, set TASK2ID, ID2TASK, refactor args
+ args = prepare_args()
+
+ # fix randomness
+ if args.seed is not None:
+ set_seed(args.seed)
+
+ # define accelerator
+ init_process_kwargs = InitProcessGroupKwargs(timeout=timedelta(seconds=args.init_timeout_seconds))
+
+ if args.distributed_type and args.distributed_type.lower() == "fsdp":
+ fsdp_plugin = FullyShardedDataParallelPlugin(
+ # state_dict_config=FullStateDictConfig(offload_to_cpu=True, rank0_only=True),
+ # optim_state_dict_config=FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=True),
+ limit_all_gathers=True,
+ sync_module_states=True,
+ use_orig_params=True,
+ cpu_offload=False,
+ )
+ accelerator = Accelerator(
+ gradient_accumulation_steps=args.gradient_accumulation_steps,
+ fsdp_plugin=fsdp_plugin,
+ dataloader_config=DataLoaderConfiguration(use_seedable_sampler=True),
+ kwargs_handlers=[init_process_kwargs],
+ )
+ else:
+ accelerator = Accelerator(
+ gradient_accumulation_steps=args.gradient_accumulation_steps,
+ dataloader_config=DataLoaderConfiguration(use_seedable_sampler=True),
+ kwargs_handlers=[init_process_kwargs],
+ )
+
+ # print key infos
+ accelerator.print("In mft_accelerate.py, sys path:", sys.path)
+ accelerator.print(f"transformers.__version__: {transformers.__version__}")
+
+ # get world_size
+ args.world_size = accelerator.num_processes
+
+ # backup args
+ pprint_args(args, accelerator)
+ if accelerator.is_main_process:
+ if not os.path.exists(args.output_dir):
+ os.makedirs(args.output_dir)
+ with open(os.path.join(args.output_dir, "args.json"), "w") as f:
+ json.dump(args.dict(), f, indent=2)
+
+ # deal with autoresume, args.resume_from_checkpoint prior to auto_resume from latest
+ latest = None
+ if os.path.exists(os.path.join(args.output_dir, "latest")):
+ with open(os.path.join(args.output_dir, "latest"), "r") as fl:
+ latest = json.load(fl)
+ accelerator.print(f"[INFO] Existing latest: {latest}")
+
+ if args.auto_resume and args.resume_from_checkpoint is None and latest:
+ args.resume_from_checkpoint = latest["latest_ckpt"]
+
+ # logger
+ logging.basicConfig(
+ format="[%(asctime)s][%(levelname)s][%(name)s]%(message)s",
+ datefmt="%m/%d/%Y %H:%M:%S",
+ level=logging.INFO,
+ )
+ logger.info(accelerator.state, main_process_only=False)
+ if accelerator.is_local_main_process:
+ datasets.utils.logging.set_verbosity_warning()
+ transformers.utils.logging.set_verbosity_info()
+ # compile Cpp helper
+ compile_helper()
+ time.sleep(10)
+ else:
+ datasets.utils.logging.set_verbosity_error()
+ transformers.utils.logging.set_verbosity_error()
+
+ # get global_rank and local rank for current process
+ global_rank = accelerator.process_index
+ local_rank = accelerator.local_process_index
+ print(f"world_size: {args.world_size}, global_rank: {global_rank}, local_rank: {local_rank}")
+
+ # TASK2ID, ID2TASK
+ # generate_task_id(args.data_paths)
+
+ # multi task blendable dataset(sharded)
+ if args.load_raw_dataset:
+ print_rank_0("> load raw jsonl dataset")
+ train_dataset, valid_dataset = load_dataset_from_jsonl(
+ args=args, shard_data=True, world_size=args.world_size, global_rank=global_rank, local_rank=local_rank
+ )
+ else:
+ print_rank_0("> load tokenized bin dataset, refer to gpt_neox indexed dataset")
+ train_dataset, valid_dataset, _ = load_dataset_from_bin(args=args)
+
+ t1 = time.time()
+ logger.info(f"dataset loading time: {t1 - t0:.4f}")
+
+ # cuda memory
+ free_in_GB = int(torch.cuda.mem_get_info()[0] / 1024**3)
+ max_memory = f"{free_in_GB - 2}GB"
+ n_gpus = torch.cuda.device_count()
+ max_memory = {i: max_memory for i in range(n_gpus)}
+ accelerator.print("max memory: ", max_memory, n_gpus)
+
+ # # 是否要加入新的special tokens
+ # num_added_toks = tokenizer.tokenizer.add_special_tokens(["", ""])
+ # accelerator.print("We have added", num_added_toks, "tokens")
+ # accelerator.print(f"role marker tokens {tokenizer.convert_tokens_to_ids('')} {tokenizer.convert_tokens_to_ids('')}, resized tokenizer_size: {len(tokenizer)}")
+
+ # creating model
+ ModelClass = MODEL_TYPES[args.model_type]
+ if args.model_type in SUPPORT_IN_TRANSFORMERS:
+ accelerator.print(f"[INFO] Model Type {args.model_type} is supported by Transformers")
+ model = ModelClass.from_pretrained(
+ args.pretrained_model_path,
+ attn_implementation=args.attn_implementation,
+ torch_dtype=torch.bfloat16,
+ )
+ else:
+ accelerator.print(f"[INFO] Model Type {args.model_type} is supported in our local model dir for remote code")
+ model = ModelClass.from_pretrained(
+ args.pretrained_model_path,
+ torch_dtype=torch.bfloat16,
+ )
+
+ # build a tokenizer for possible resizing or saving
+ tokenizer = build_tokenizer(args)
+ # Note: resize_token_embeddings expects to receive the full size of the new vocabulary,
+ # i.e. the length of the tokenizer.
+ # 如果新增special tokens, 需要resize input embedding 和output embedding
+ # model.resize_token_embeddings(len(tokenizer), pad_to_multiple_of=32)
+
+ model.gradient_checkpointing_enable()
+
+ if args.saving_limit is None or not isinstance(args.saving_limit, int) or args.saving_limit < 1:
+ # saving_limit is set automatically if needed
+ args.saving_limit = 2
+ accelerator.print(
+ "[WARNING]saving_limit must be a integer greater than 1 in Full-Parameters Training, we set it to 2"
+ )
+
+ t2 = time.time()
+ if accelerator.is_main_process:
+ logging.info(f"model loading time: {t2 - t1:.4f}")
+
+ model.config.use_cache = False # silence the warnings. Please re-enable for inference!
+ if hasattr(model.config, "use_logn_attn"):
+ model.config.use_logn_attn = False # special for qwen model
+ # load balance for moe training
+ if hasattr(model.config, "output_router_logits"):
+ model.config.output_router_logits = True
+ model_config = model.config
+ accelerator.print(model.config)
+
+ # dataloader
+ train_dataloader = DataLoader(
+ train_dataset,
+ shuffle=True,
+ collate_fn=DataCollatorForMFTDataset(args),
+ batch_size=args.per_device_train_batch_size,
+ pin_memory=True,
+ drop_last=True,
+ )
+ if valid_dataset:
+ valid_dataloader = DataLoader(
+ valid_dataset,
+ collate_fn=DataCollatorForMFTDataset(args),
+ batch_size=args.per_device_eval_batch_size,
+ pin_memory=True,
+ drop_last=True,
+ )
+ else:
+ valid_dataloader = None
+
+ # optimizer
+ if accelerator.distributed_type == DistributedType.DEEPSPEED:
+ accelerator.print("DISTRIBUTED TRAINING USING DEEPSPEED")
+ # from deepspeed.ops.adam import FusedAdam as Adam
+ # adam_optimizer = Adam
+ adam_optimizer = torch.optim.AdamW
+ elif accelerator.distributed_type == DistributedType.FSDP:
+ accelerator.print("DISTRIBUTED TRAINING USING FSDP")
+ model = accelerator.prepare(model)
+ adam_optimizer = torch.optim.AdamW
+ else:
+ raise ValueError("Only support DeepSpeed and FSDP")
+
+ optimizer = adam_optimizer(
+ model.parameters(),
+ weight_decay=args.weight_decay,
+ lr=args.learning_rate,
+ betas=(0.9, 0.999),
+ )
+ # for group in optimizer.param_groups:
+ # group.setdefault("initial_lr", group["lr"])
+
+ # Scheduler and math around the number of training steps.
+ overrode_max_train_steps = False
+ num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+ if args.max_train_steps is None:
+ args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
+ overrode_max_train_steps = True
+ if isinstance(args.num_warmup_steps, float) and args.num_warmup_steps < 1.0:
+ args.num_warmup_steps = int(args.max_train_steps * args.num_warmup_steps) // accelerator.num_processes
+ accelerator.print(f"num_warmup_steps: {args.num_warmup_steps}")
+ lr_scheduler = get_scheduler(
+ name=args.lr_scheduler_type,
+ optimizer=optimizer,
+ num_warmup_steps=args.num_warmup_steps * args.gradient_accumulation_steps * accelerator.num_processes,
+ num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
+ # scheduler_specific_kwargs={"last_epoch": scheduler_last_ep}
+ )
+ # prepare all
+ if accelerator.distributed_type == DistributedType.DEEPSPEED:
+ if valid_dataloader:
+ (model, train_dataloader, valid_dataloader, optimizer, lr_scheduler) = accelerator.prepare(
+ model, train_dataloader, valid_dataloader, optimizer, lr_scheduler
+ )
+ else:
+ (model, train_dataloader, optimizer, lr_scheduler) = accelerator.prepare(
+ model, train_dataloader, optimizer, lr_scheduler
+ )
+
+ # prepare all except model, which is prepared before
+ elif accelerator.distributed_type == DistributedType.FSDP:
+ if valid_dataloader:
+ (optimizer, train_dataloader, valid_dataloader, lr_scheduler) = accelerator.prepare(
+ optimizer, train_dataloader, valid_dataloader, lr_scheduler
+ )
+ else:
+ (optimizer, train_dataloader, lr_scheduler) = accelerator.prepare(
+ optimizer, train_dataloader, lr_scheduler
+ )
+ print(model.device)
+ accelerator.print(model)
+ # accelerator.print(model.config)
+
+ # Recalculate our total training steps as the size of the training dataloader may have changed.
+ num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+ if overrode_max_train_steps:
+ args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
+ # Afterward we recalculate our number of training epochs
+ args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+
+ # zero 3 flag
+ is_ds_zero_3 = False
+ if getattr(accelerator.state, "deepspeed_plugin", None):
+ is_ds_zero_3 = accelerator.state.deepspeed_plugin.zero_stage == 3
+ accelerator.print(f"DEEPSPEED plugin: {accelerator.state.deepspeed_plugin}")
+ elif getattr(accelerator.state, "fsdp_plugin", None):
+ accelerator.print(f"FSDP plugin: {accelerator.state.fsdp_plugin}")
+
+ trainer = MptTrainer(
+ accelerator=accelerator,
+ model=model,
+ model_config=model_config,
+ train_dataloader=train_dataloader,
+ valid_dataloader=valid_dataloader,
+ optimizer=optimizer,
+ lr_scheduler=lr_scheduler,
+ tokenizer=tokenizer,
+ num_update_steps_per_epoch=num_update_steps_per_epoch,
+ total_train_dataset_size=len(train_dataset),
+ args=args,
+ )
+ trainer.accelerate_train()
+
+
+if __name__ == "__main__":
+ main()
diff --git a/mftcoder_accelerate/src/mpt/mpt_arguments.py b/mftcoder_accelerate/src/mpt/mpt_arguments.py
new file mode 100644
index 0000000..8045421
--- /dev/null
+++ b/mftcoder_accelerate/src/mpt/mpt_arguments.py
@@ -0,0 +1,161 @@
+"""
+# @author Chaoyu Chen
+# @date 2024/6/1
+
+MPT training arguments
+"""
+
+from dataclasses import dataclass, asdict
+from typing import List, Union
+
+
+@dataclass
+class MptTrainArgs:
+ # train data paths on shared FS
+ data_paths: Union[str, List[str]]
+
+ # output dir for saving adaptors in peft or full ckpts in full-parameter training
+ output_dir: str
+
+ # tensorboard dir for saving tensorboard logs
+ tb_dir: str
+
+ # pretrained_model_path, on which is the model you want to train
+ pretrained_model_path: str
+
+ # model type of pretrained_model_path, support llama|qwen|starcoder|baichuan|chatglm2
+ model_type: str
+
+ # load from raw jsonl file or tokenized binary file
+ load_raw_dataset: bool = True
+
+ # weights of loss calculation for each task, None means equal weights
+ task_weights: Union[None, str] = None
+
+ # weights of data sampling, leave it None
+ data_weights: Union[None, str] = None
+
+ # hf loading model low_cpu_mem_usage
+ low_cpu_mem_usage: bool = True
+
+ # train/valid/test split
+ data_split: str = "98,2,0"
+
+ # padding or pack or concat
+ padding_mode: str = "padding"
+
+ # sft or sst
+ tokenize_mode: str = "sft"
+
+ # case3 or case4
+ weighted_loss_mode: str = "case3"
+
+ # mircro train batch size
+ per_device_train_batch_size: int = 8
+
+ # micro eval batch size, always same as micro train batch size
+ per_device_eval_batch_size: int = 8
+
+ # HF AutoTokenizer is supported, maybe more types
+ tokenizer_type: str = "AutoTokenizer"
+
+ # initial lr
+ learning_rate: float = 5e-5
+
+ # minimum lr
+ min_lr: float = 5e-6
+
+ # weight decay
+ weight_decay: float = 0.01
+
+ # gradient_accumulation_steps
+ gradient_accumulation_steps: int = 1
+
+ # lr_scheduler_type
+ lr_scheduler_type: str = "cosine"
+
+ # num_warmup_steps
+ num_warmup_steps: Union[int, float] = 0.05
+
+ # num_train_epochs
+ num_train_epochs: int = 4
+
+ # seed for reproducing
+ seed: int = 1234
+
+ # seq_length, context length
+ seq_length: int = 4096
+
+ # path of adaptor which is resumed from, None for not resuming training
+ resume_from_checkpoint: Union[None, str] = None
+
+ # auto resume from latest ckpt if job restarted
+ auto_resume: bool = True
+
+ # num of steps for logging training loss
+ log_interval: int = 10
+
+ # num of steps for saving ckpt
+ checkpointing_steps: int = 100
+
+ # num of steps for evaluation(eval_loss), better same as checkpointing steps
+ evaluation_steps: int = 100
+
+ # max train steps, if None, depends on num_train_epochs
+ max_train_steps: Union[None, int] = None
+
+ # if checkpointing every epoch, maybe True in sst
+ epoch_checkpointing: bool = False
+
+ # save transformers model(safetensors)
+ save_transformers_model: bool = False
+
+ # shuffle before train/valid split
+ shuffle_before_split: bool = True
+
+ # DDP random sampler
+ use_random_sampler: bool = True
+
+ # if early stop when eval loss is not converging in the past early_stopping_stall_num evaluation point
+ early_stopping: bool = True
+ early_stopping_stall_num: int = 5
+
+ # limit num for saving ckpts, None for no limits. Used for full-parameter training to avoid exceeding disk quota.
+ saving_limit: Union[None, int] = None
+
+ # if dynamic padding
+ use_dynamic_padding: bool = True
+
+ # warm-up steps for CoBa, recommand the number of valid batches
+ coba_warmup_steps: int = 100
+ # history length of sample valid loss used to fit the slope curve in CoBa, recommand [2*coba_warmup_steps,5*coba_warmup_steps]
+ coba_history_length: int = 200
+ # temperature for divergence factor in CoBa
+ coba_tau: int = 5
+ # iteration interval of update per task train weight in CoBa
+ coba_update_interval: int = 1
+ # the number of mini valid batches sampled at each updated iteration interval
+ coba_sample_valid_num: int = 1
+
+ # ATTENTION_CLASSES = { "eager": Normal Attention, "flash_attention_2": FlashAttention2}
+ attn_implementation: str = "flash_attention_2"
+
+ # role markers, which are prompt template before each role: system, user and assistant
+ # role_markers: {"system": "### System:\n", "user": "### Instruction:\n", "assistant": "### Response:\n"}
+ role_markers: Union[None, dict] = None
+
+ distributed_type: Union[None, str] = None
+
+ init_timeout_seconds: Union[None, int] = 3600
+
+ # legacy, leave them
+ use_xformers: bool = True
+ trust_remote_code: bool = True
+ weight_by_num_documents: bool = True
+ make_vocab_size_divisible_by: int = 32
+ model_parallel_size: int = 1
+ use_slow_tokenizer: bool = False
+ world_size: int = 8
+
+ def dict(self):
+ return {k: str(v) for k, v in asdict(self).items()}
diff --git a/mftcoder_accelerate/src/mpt/mpt_trainer.py b/mftcoder_accelerate/src/mpt/mpt_trainer.py
new file mode 100644
index 0000000..b5e2da8
--- /dev/null
+++ b/mftcoder_accelerate/src/mpt/mpt_trainer.py
@@ -0,0 +1,606 @@
+"""
+# @author qumu
+# @date 2024/6/6
+# @module mpt_trainer.py
+
+MPT/MCT/MFT Full-parameter Trainer
+"""
+
+import gc
+import os
+import sys
+import threading
+import argparse
+import math
+import logging
+import json
+import time
+import transformers
+import numpy as np
+import psutil
+import shutil
+import torch
+from torch import nn
+from torch.utils.tensorboard import SummaryWriter
+from typing import List, Optional, Tuple, Union
+from tqdm.auto import tqdm
+from accelerate.logging import get_logger
+from accelerate import Accelerator
+from transformers import set_seed
+
+# sys.path.append("..")
+from utils.common_utils import generate_task_id, TASK2ID, ID2TASK
+from utils.loss_utils import loss_func_mft, CoBaStatus, load_balancing_loss_func
+
+logger = get_logger(__name__)
+
+
+def copy_tokenizer_files(mode_path: str, files_list: List[str], save_path: str):
+ # create path if not exist
+ if not os.path.exists(save_path):
+ os.makedirs(save_path)
+
+ # copy each file in files_list to save_path
+ for filename in files_list:
+ src_file = os.path.join(mode_path, filename)
+
+ # copy only if src exists
+ if os.path.exists(src_file):
+ dest_file = os.path.join(save_path, filename)
+
+ # copy
+ shutil.copy(src_file, dest_file)
+ print(f"Copied {filename} to {save_path}")
+ else:
+ print(f"File {filename} does not exist in {mode_path}")
+
+
+def check_existing_ckpts(output_dir):
+ prefix = "step_"
+
+ if not os.path.exists(output_dir):
+ return []
+ # list all files and dirs
+ contents = os.listdir(output_dir)
+
+ # find dirs starts with "step_"
+ matching_folders = [
+ folder for folder in contents if os.path.isdir(os.path.join(output_dir, folder)) and folder.startswith(prefix)
+ ]
+
+ return matching_folders
+
+
+def extract_epochs_and_steps(path, num_update_steps_per_epoch, gradient_accumulation_steps):
+ """
+ extract starting_epoch, completed_steps, resume_step of train_dataloader for resumed training
+ """
+ # Extract `epoch_{i}` or `step_{i}`
+ training_difference = os.path.splitext(path)[0]
+
+ if "epoch" in training_difference:
+ starting_epoch = int(training_difference.replace("epoch_", ""))
+ resume_step = None
+ completed_steps = starting_epoch * num_update_steps_per_epoch
+ logger.info(f"Resume from exact Epoch {starting_epoch}: completed_steps {completed_steps}")
+ else:
+ # need to multiply `gradient_accumulation_steps` to reflect real steps
+ completed_steps = int(training_difference.replace("step_", ""))
+ starting_epoch = completed_steps // num_update_steps_per_epoch
+ resume_step = (completed_steps % num_update_steps_per_epoch) * gradient_accumulation_steps
+ logger.info(f"Resume from Epoch {starting_epoch} + step {resume_step}: completed_steps {completed_steps}")
+
+ return starting_epoch, completed_steps, resume_step
+
+
+def write_tensorboard(summary_writer: SummaryWriter, log_dict: dict, completed_steps):
+ for key, value in log_dict.items():
+ summary_writer.add_scalar(f"{key}", value, completed_steps)
+
+
+def delete_ckpts_over_limits(output_dir, saving_limit, best_step):
+ """delete ckpts more than saving_limits except for the best_step ckpt"""
+ existing_ckpts = check_existing_ckpts(output_dir)
+ logger.info(f"Existing step ckpts folders: {existing_ckpts}, best step ckpt: step_{best_step}")
+ # sorted only step num ascendingly
+ ckpt_steps = sorted([int(ckpt.replace("step_", "")) for ckpt in existing_ckpts])
+ # delete the oldest steps except for the best step at present
+ if len(ckpt_steps) > saving_limit:
+ deletable_steps = [ckpt_step for ckpt_step in ckpt_steps if ckpt_step != best_step]
+ # print(deletable_steps[:len(ckpt_steps) - saving_limit])
+ for del_step in deletable_steps[: len(ckpt_steps) - saving_limit]:
+ shutil.rmtree(os.path.join(output_dir, f"step_{del_step}"))
+ logger.info(f"Removed ckpt step_{del_step}")
+
+
+class MptTrainer:
+ """
+ Multitask Pre-train/Continue-train Trainer with Full-parameters training.
+ """
+
+ def __init__(
+ self,
+ accelerator: Accelerator,
+ model,
+ model_config,
+ train_dataloader,
+ valid_dataloader,
+ optimizer,
+ lr_scheduler,
+ tokenizer,
+ num_update_steps_per_epoch,
+ total_train_dataset_size,
+ args,
+ ):
+ self.accelerator = accelerator
+ self.model = model
+ # hf model config
+ self.model_config = model_config
+ self.train_dataloader = train_dataloader
+ self.valid_dataloader = valid_dataloader
+ self.optimizer = optimizer
+ self.lr_scheduler = lr_scheduler
+ self.tokenizer = tokenizer
+ self.num_update_steps_per_epoch = num_update_steps_per_epoch
+ self.total_train_dataset_size = total_train_dataset_size
+ # training arguments
+ self.args = args
+ # tensorboard writer
+ self.summary_writer = SummaryWriter(log_dir=args.tb_dir)
+
+ def print(self, msg: str):
+ """
+ accelerator print, default on main process
+ Args:
+ msg:
+
+ Returns:
+
+ """
+ self.accelerator.print(msg)
+
+ def touch(self, batch, num_tokens=10):
+ """touch first and last tokens and labels for debugging usage"""
+ self.print(
+ f"step 1 batch shape: {batch['input_ids'].shape},\n"
+ f"last {num_tokens} labels: {batch['labels'][:, -num_tokens:]}"
+ f"last {num_tokens} loss mask: {batch['loss_mask'][:, -num_tokens:]}"
+ )
+ self.print(f"first {num_tokens} input_ids and loss_mask")
+ for pt in range(1):
+ self.print(f"{batch['input_ids'][:, num_tokens * pt: num_tokens * pt + num_tokens]}")
+ self.print(f"{batch['loss_mask'][:, num_tokens * pt: num_tokens * pt + num_tokens]}")
+
+ @staticmethod
+ def format_tensor(tensor, n):
+ return list(map(lambda x: round(x, n), tensor.tolist()))
+
+ def accelerate_saving_states(self, output_dir: str, completed_steps: int):
+ """
+ Saving lora adaptor or full checkpoint using accelerator
+ Args:
+ output_dir: exact dir for saving ckpt
+ completed_steps:
+
+ Returns:
+
+ """
+ self.accelerator.wait_for_everyone()
+ logger.info(f"[CHECKPOINT] Saving checkpoint states")
+ self.accelerator.save_state(output_dir)
+ self.accelerator.wait_for_everyone()
+
+ # save safetensors for direct inference if needed
+ if self.args.save_transformers_model:
+ logger.info(f"[CHECKPOINT] Saving transformers(hf) model", main_process_only=True)
+ unwrapped_model = self.accelerator.unwrap_model(self.model)
+ # self.print(f"unwrapped model type {type(unwrapped_model)}")
+ unwrapped_model.save_pretrained(
+ output_dir,
+ is_main_process=self.accelerator.is_main_process,
+ save_function=self.accelerator.save,
+ state_dict=self.accelerator.get_state_dict(self.model),
+ )
+ self.accelerator.wait_for_everyone()
+
+ # tokenizer saving and bug dummy ckpt cleaning.
+ if self.accelerator.is_main_process:
+ if self.args.model_type.lower() == "deepseek":
+ copy_tokenizer_files(
+ self.args.pretrained_model_path, ["tokenizer.json", "tokenizer_config.json"], output_dir
+ )
+ else:
+ self.tokenizer.save_pretrained(output_dir)
+
+ sf = os.path.join(output_dir, "model.safetensors")
+ index_file = os.path.join(output_dir, "model.safetensors.index.json")
+ if os.path.isfile(sf) and os.path.isfile(index_file):
+ self.print(f"Remove bug dummy ckpt {sf}")
+ os.remove(sf)
+
+ # save latest info
+ if self.accelerator.is_main_process:
+ latest = {
+ "latest_ckpt": output_dir,
+ "lr": self.optimizer.param_groups[0]["lr"],
+ }
+ with open(os.path.join(self.args.output_dir, "latest"), "w") as f:
+ json.dump(latest, f, indent=2)
+
+ logger.info(
+ f"[CHECKPOINT][complete_steps={completed_steps}], states {output_dir} saved, latest: {latest}",
+ main_process_only=True,
+ )
+ self.accelerator.wait_for_everyone()
+
+ def accelerate_monitor(
+ self,
+ reduce_loss,
+ reduce_task_loss,
+ reduce_task_exist,
+ completed_steps,
+ coba_status=None,
+ ):
+ """
+ gather reduce_loss and reduce_task_loss from all N devices.
+ train logging and tensorboarding.
+ """
+ # gather reduce_loss and reduce_task_loss from all N devices
+ reduce_losses = self.accelerator.gather(reduce_loss).detach().float()
+ reduce_task_losses = self.accelerator.gather(reduce_task_loss).reshape(-1, len(ID2TASK))
+ reduce_task_exists = self.accelerator.gather(reduce_task_exist).reshape(-1, len(ID2TASK))
+
+ # get train loss and per-task train loss
+ train_loss = torch.mean(reduce_losses) / (self.args.log_interval * self.args.gradient_accumulation_steps)
+ # train_task_loss = torch.mean(reduce_task_losses, dim=0) / (self.args.log_interval * self.args.gradient_accumulation_steps)
+ train_task_loss = torch.sum(reduce_task_losses, dim=0) / torch.sum(reduce_task_exists, dim=0)
+
+ # logging and writing tensorboard
+ logger.info(
+ f"[TRAIN][complete_steps={completed_steps}][train_loss={train_loss:.6f}]"
+ f"[train_task_loss={self.format_tensor(train_task_loss, 4)}]"
+ f"[gather shape={list(reduce_losses.shape)}]"
+ f"[lr={self.lr_scheduler.get_lr()[0]:.4e}, {self.optimizer.param_groups[0]['lr']:.4e}]",
+ main_process_only=True,
+ )
+ if coba_status is not None:
+ if completed_steps > coba_status.coba_warmup_steps:
+ coba_status.log_per_task_weight = coba_status.log_per_task_weight / torch.sum(
+ coba_status.log_per_task_weight
+ )
+ else:
+ coba_status.log_per_task_weight = torch.ones(len(ID2TASK)) / len(ID2TASK)
+ logger.info(
+ f"[TRAIN][per_task_train_weight={coba_status.log_per_task_weight}]", main_process_only=True
+ )
+ train_log_dict = {"Loss/train": train_loss}
+ for i in range(len(ID2TASK)):
+ train_log_dict[f"{ID2TASK[i]}_loss/train"] = train_task_loss[i]
+ if coba_status is not None:
+ train_log_dict[f"{ID2TASK[i]}_coba_weight/train"] = coba_status.log_per_task_weight[i].item()
+
+ if self.accelerator.is_main_process:
+ write_tensorboard(self.summary_writer, train_log_dict, completed_steps)
+
+ if coba_status is not None:
+ coba_status.log_per_task_weight = torch.zeros(len(ID2TASK))
+
+ def accelerate_evaluate(
+ self,
+ completed_steps,
+ step,
+ min_eval_loss,
+ stall_num,
+ best_step,
+ ):
+ """
+ evaluate the model at current completed_steps on valid_dataloader and gather eval_loss on all devices.
+ eval logging and tensorboarding.
+ """
+ losses = []
+ accumulated_task_loss = torch.zeros(len(ID2TASK)).to(self.model.device)
+ accumulated_task_exist = torch.zeros(len(ID2TASK)).to(self.model.device)
+ for valid_step, valid_batch in enumerate(self.valid_dataloader):
+ with torch.no_grad():
+ outputs = self.model(
+ input_ids=valid_batch["input_ids"],
+ attention_mask=valid_batch["attention_mask"],
+ position_ids=valid_batch["position_ids"],
+ return_dict=True,
+ )
+
+ loss, task_loss, _ = loss_func_mft(
+ outputs=outputs,
+ labels=valid_batch["labels"],
+ task_mask=valid_batch["task_mask"],
+ task_id=valid_batch["task_id"],
+ weighted_loss_mode=self.args.weighted_loss_mode,
+ loss_mask=valid_batch["loss_mask"],
+ task_weights=self.args.task_weights,
+ )
+
+ losses.append(self.accelerator.gather(loss.repeat(self.args.per_device_eval_batch_size)))
+ accumulated_task_loss += task_loss.detach().float()
+ accumulated_task_exist += (task_loss != 0.0).detach().float()
+
+ self.accelerator.wait_for_everyone()
+ valid_batch_num = len(losses)
+ gathered_size = losses[0].shape
+ losses = torch.cat(losses)
+ # task_losses = torch.cat(task_losses).reshape(-1, len(ID2TASK))
+ task_losses = self.accelerator.gather(accumulated_task_loss).reshape(-1, len(ID2TASK))
+ task_exists = self.accelerator.gather(accumulated_task_exist).reshape(-1, len(ID2TASK))
+
+ try:
+ eval_loss = torch.mean(losses)
+ # eval_task_loss = torch.mean(task_losses, dim=0) / valid_batch_num
+ eval_task_loss = torch.sum(task_losses, dim=0) / torch.sum(task_exists, dim=0)
+ if eval_loss <= min_eval_loss:
+ min_eval_loss = eval_loss
+ stall_num = 0
+ best_step = completed_steps
+ else:
+ stall_num += 1
+ perplexity = math.exp(eval_loss)
+ except OverflowError:
+ perplexity = float("inf")
+
+ logger.info(
+ f"[EVAL][completed_steps={completed_steps}]"
+ f"[eval_loss={eval_loss:.6f}][eval_task_loss={self.format_tensor(eval_task_loss, 4)}]"
+ f"[perplexity={perplexity:.4f}][valid_batch_num={valid_batch_num}]"
+ f"[gather_size={list(gathered_size)}]",
+ main_process_only=True,
+ )
+ eval_log_dict = {
+ "Loss/valid": eval_loss,
+ "Perplexity/valid": perplexity,
+ "Epochs": round(completed_steps / self.num_update_steps_per_epoch, 2),
+ }
+ for i in range(len(ID2TASK)):
+ eval_log_dict[f"{ID2TASK[i]}_loss/valid"] = eval_task_loss[i]
+
+ if self.accelerator.is_main_process:
+ write_tensorboard(self.summary_writer, eval_log_dict, completed_steps)
+
+ return eval_loss, eval_task_loss, min_eval_loss, stall_num, best_step
+
+ def accelerate_train(self):
+ # Train!
+ if self.args.seed is not None:
+ set_seed(self.args.seed)
+
+ global_batch_size = (
+ self.args.per_device_train_batch_size
+ * self.accelerator.num_processes
+ * self.args.gradient_accumulation_steps
+ )
+ logger.info("************************************** Running training ****************************************")
+ logger.info(f" Num examples = {self.total_train_dataset_size}")
+ logger.info(f" Num Epochs = {self.args.num_train_epochs}")
+ logger.info(f" Instantaneous batch size per device = {self.args.per_device_train_batch_size}")
+ logger.info(f" Total global train batch size (w. parallel, distributed & accumulation) = {global_batch_size}")
+ logger.info(f" Gradient Accumulation steps = {self.args.gradient_accumulation_steps}")
+ logger.info(f" Total optimization(update/completed) steps = {self.args.max_train_steps}")
+ logger.info(f" Complete/optimize steps per Epoch = {self.args.max_train_steps // self.args.num_train_epochs}")
+ logger.info("************************************************************************************************")
+
+ # Only show the progress bar once on each machine.
+ progress_bar = tqdm(range(self.args.max_train_steps), disable=not self.accelerator.is_local_main_process)
+
+ # set starting_epoch, completed_steps and resume_step of train_dataloader
+ completed_steps = 0
+ starting_epoch = 0
+ resume_step = None
+
+ if self.args.resume_from_checkpoint:
+ self.accelerator.load_state(self.args.resume_from_checkpoint)
+ self.accelerator.print(f"Resumed from checkpoint: {self.args.resume_from_checkpoint}")
+ path = os.path.basename(self.args.resume_from_checkpoint)
+ starting_epoch, completed_steps, resume_step = extract_epochs_and_steps(
+ path, self.num_update_steps_per_epoch, self.args.gradient_accumulation_steps
+ )
+
+ # update the progress_bar if load from checkpoint
+ progress_bar.update(completed_steps)
+
+ # monitor minimum eval_loss, stalling num, and best_step
+ min_eval_loss = float("inf")
+ stall_num = 0
+ best_step = None
+
+ # monitor train loss
+ reduce_loss = torch.tensor(0.0).to(self.model.device)
+ reduce_aux_loss = torch.tensor(0.0).to(self.model.device)
+ reduce_task_loss = torch.zeros(len(ID2TASK)).to(self.model.device)
+ reduce_task_exist = torch.zeros(len(ID2TASK)).to(self.model.device)
+ per_task_weight = self.args.task_weights
+
+ if self.args.weighted_loss_mode == "coba":
+ self.model.eval()
+ eval_loss, eval_task_loss, _, _, _ = self.accelerate_evaluate(
+ completed_steps,
+ 0,
+ min_eval_loss,
+ stall_num,
+ best_step,
+ )
+ self.model.train()
+ coba_status = CoBaStatus(
+ self.args.coba_warmup_steps,
+ self.args.coba_history_length,
+ self.args.coba_tau,
+ self.args.coba_update_interval,
+ self.args.coba_sample_valid_num,
+ self.valid_dataloader,
+ )
+ coba_status.valid_task_loss_begining = eval_task_loss.clone().to(self.model.device)
+ coba_status.sample_valid_batch(self.model, completed_steps)
+ logger.info(f"valid_task_loss: {coba_status.valid_task_loss_accumulated}", main_process_only=True)
+ else:
+ coba_status = None
+
+ # Training Loop!
+ for epoch in range(starting_epoch, self.args.num_train_epochs):
+ # set_epoch
+ # self.train_dataloader.set_epoch(epoch)
+
+ # if we early stop by some ckpts not converging
+ if self.args.early_stopping and stall_num == self.args.early_stopping_stall_num:
+ break
+
+ if self.args.resume_from_checkpoint and epoch == starting_epoch and resume_step is not None:
+ # We skip the first `n` batches in the dataloader when resuming from a checkpoint
+ active_dataloader = self.accelerator.skip_first_batches(self.train_dataloader, resume_step)
+ else:
+ active_dataloader = self.train_dataloader
+ tail_num = len(active_dataloader) - len(active_dataloader) % self.args.gradient_accumulation_steps
+ print(f"length of dataloader: {len(active_dataloader)}")
+
+ self.model.train()
+ # Inner Loop!
+ for step, batch in enumerate(active_dataloader):
+ if step == tail_num:
+ break
+ with self.accelerator.accumulate(self.model):
+ if step == 0:
+ self.touch(batch, num_tokens=10)
+ # forward
+ outputs = self.model(
+ input_ids=batch["input_ids"],
+ attention_mask=batch["attention_mask"],
+ position_ids=batch["position_ids"],
+ return_dict=True,
+ )
+
+ if (
+ self.args.weighted_loss_mode == "coba"
+ and self.accelerator.sync_gradients
+ and completed_steps % self.args.coba_update_interval == 0
+ and completed_steps >= self.args.coba_warmup_steps
+ ):
+ with torch.no_grad():
+ per_task_weight = coba_status.compute_per_task_weight(completed_steps=completed_steps)
+ coba_status.log_per_task_weight += per_task_weight
+ # logger.info(f'per_task_weight: {per_task_weight}', main_process_only=True)
+
+ # loss
+ loss, task_loss, _ = loss_func_mft(
+ outputs=outputs,
+ labels=batch["labels"],
+ task_mask=batch["task_mask"],
+ task_id=batch["task_id"],
+ weighted_loss_mode=self.args.weighted_loss_mode,
+ loss_mask=batch["loss_mask"],
+ task_weights=per_task_weight,
+ )
+
+ # accelerator.print(len(outputs.router_logits), outputs.router_logits[0], outputs.router_logits[-1])
+ # accelerator.print(batch['attention_mask'].shape, batch['attention_mask'])
+ aux_loss = None
+ if hasattr(self.model_config, "output_router_logits") and self.model_config.output_router_logits:
+ if hasattr(self.model_config, "num_local_experts"):
+ num_experts = self.model_config.num_local_experts
+ elif hasattr(self.model_config, "num_experts"):
+ num_experts = self.model_config.num_experts
+ else:
+ raise ValueError("model has no attribute num_local_experts or num_experts")
+ aux_loss = load_balancing_loss_func(
+ outputs.router_logits,
+ num_experts,
+ self.model_config.num_experts_per_tok,
+ batch["attention_mask"],
+ )
+ aux_loss = self.model_config.router_aux_loss_coef * aux_loss.to(loss.device)
+ loss += aux_loss # make sure to reside in the same device
+
+ # backward
+ self.accelerator.backward(loss)
+ # print(self.lr_scheduler.state_dict(), self.accelerator.process_index)
+ # update(sync_gradients)
+ self.optimizer.step()
+ self.lr_scheduler.step()
+ self.optimizer.zero_grad()
+ # support args.min_lr
+ if self.optimizer.param_groups[0]["lr"] <= self.args.min_lr:
+ self.optimizer.param_groups[0]["lr"] = self.args.min_lr
+
+ # accumulate resuce_loss and reduce_task_loss in a log_interval
+ if not torch.isnan(loss):
+ reduce_loss += loss.detach().float()
+ if aux_loss and not torch.isnan(aux_loss):
+ reduce_aux_loss += aux_loss.detach().float()
+ # self.print("task loss devices: ", reduce_task_loss.device, task_loss.device)
+ reduce_task_loss += task_loss.detach().float()
+ reduce_task_exist += (task_loss != 0).detach().float()
+
+ # If the accelerator has performed an optimization step behind the scenes, thus a completed_step done.
+ if self.accelerator.sync_gradients:
+ if (
+ self.args.weighted_loss_mode == "coba"
+ and completed_steps % self.args.coba_update_interval == 0
+ and completed_steps >= 1
+ ):
+ coba_status.sample_valid_batch(self.model, completed_steps)
+ # logger.info(f"valid_task_loss: {coba_status.valid_task_loss_accumulated}", main_process_only=True)
+
+ # progress_bar.update(1)
+ completed_steps += 1
+ # monitoring training process and logging and tensorboarding
+ if completed_steps % self.args.log_interval == 0:
+ progress_bar.update(self.args.log_interval)
+ if reduce_aux_loss > 0.0:
+ self.print(f"[INFO] aux_loss: {reduce_aux_loss/self.args.log_interval}")
+ self.accelerate_monitor(
+ reduce_loss,
+ reduce_task_loss,
+ reduce_task_exist,
+ completed_steps,
+ coba_status,
+ )
+ # reset reduce_loss
+ reduce_loss = torch.tensor(0.0).to(self.model.device)
+ reduce_aux_loss = torch.tensor(0.0).to(self.model.device)
+ reduce_task_loss = torch.zeros(len(ID2TASK)).to(self.model.device)
+ reduce_task_exist = torch.zeros(len(ID2TASK)).to(self.model.device)
+
+ # steps checkpointing
+ if self.args.checkpointing_steps and completed_steps % self.args.checkpointing_steps == 0:
+ output_dir = f"step_{completed_steps}"
+ if self.args.output_dir is not None:
+ output_dir = os.path.join(self.args.output_dir, output_dir)
+ self.accelerate_saving_states(output_dir, completed_steps)
+
+ # steps evaluation
+ if completed_steps % self.args.evaluation_steps == 0 and self.valid_dataloader:
+ self.model.eval()
+ eval_loss, eval_task_loss, min_eval_loss, stall_num, best_step = self.accelerate_evaluate(
+ completed_steps,
+ step,
+ min_eval_loss,
+ stall_num,
+ best_step,
+ )
+ self.model.train()
+
+ # delete ckpts over args.saving_limit
+ if self.accelerator.is_main_process and self.args.saving_limit:
+ delete_ckpts_over_limits(self.args.output_dir, self.args.saving_limit, best_step)
+
+ # early stoppin when stalling more than args.early_stopping_stall_num
+ if self.args.early_stopping and stall_num == self.args.early_stopping_stall_num:
+ self.print(f"[WARNING] Early stopping at {completed_steps}")
+ break
+
+ if completed_steps >= self.args.max_train_steps:
+ break
+ self.accelerator.wait_for_everyone()
+
+ # epoch checkpointing
+ if self.args.epoch_checkpointing:
+ output_dir = f"epoch_{epoch + 1}"
+ if self.args.output_dir is not None:
+ output_dir = os.path.join(self.args.output_dir, output_dir)
+ self.accelerate_saving_states(output_dir, completed_steps)
+
+ self.summary_writer.close()
diff --git a/mftcoder_accelerate/src/offline_tokenization/concat_sst_bin_tokenization.py b/mftcoder_accelerate/src/offline_tokenization/concat_sst_bin_tokenization.py
new file mode 100644
index 0000000..ca4347e
--- /dev/null
+++ b/mftcoder_accelerate/src/offline_tokenization/concat_sst_bin_tokenization.py
@@ -0,0 +1,220 @@
+# -*- coding: utf-8 -*-
+
+import argparse
+import multiprocessing
+import os
+import sys
+import random
+import time
+import tqdm
+import glob
+import json
+import numpy as np
+
+
+# 将父目录的父目录加入path
+current_path = os.path.abspath(__file__)
+parent_dir = os.path.dirname(os.path.dirname(current_path))
+grandparent_dir = os.path.dirname(parent_dir)
+sys.path.append(grandparent_dir)
+
+from tokenizer import init_tokenizer
+from pack_encoder import PackSSTBinEncoder, load_tokenizer
+from data import indexed_dataset
+
+from threading import Semaphore
+from colorama import Fore
+import lm_fmt as lmd
+
+
+def yield_from_files(files: list, semaphore):
+ """
+ Iterator over input documents
+
+ :param fnames: list of filenames
+ """
+ def yielder(fname, semaphore):
+ with open(fname, 'r') as f:
+ for line in f:
+ semaphore.acquire()
+ yield json.loads(line)
+
+ for fname in files:
+ semaphore.acquire()
+ yield from yielder(fname, semaphore)
+
+def yield_from_files2(fnames: list, semaphore, sample_percent):
+ """
+ Iterator over input documents using lm_dataformat. Should be able to handle jsons / texts /
+ other compressed formats. Also filters out empty documents.
+
+ :param fnames: list of filenames
+ """
+ def yielder(fname, semaphore):
+ try:
+ sample_interval = int(1/sample_percent)
+ for f in filter(lambda x: x, lmd.Reader(fname).stream_data(key=None)):
+ rand_value = random.randint(1, sample_interval*100)
+ if rand_value % sample_interval != 0:
+ continue
+ semaphore.acquire()
+
+ #rand_value = random.randint(1, sample_interval*100)
+ #if rand_value % sample_interval != 0:
+ # yield None
+
+ yield f
+ except Exception as e:
+ print('####Exception:', e.args)
+ yield None
+
+ for fname in fnames:
+ semaphore.acquire()
+
+ yield from yielder(fname, semaphore)
+
+
+def print_example_doc(input_ids, tokenizer):
+ print(Fore.YELLOW + f'INPUT IDS len: {len(input_ids)}')
+ print(Fore.BLUE + f'INPUT IDS:\n {input_ids}\n\n')
+
+ print(Fore.RED + f'DETOKENIZED INPUT:\n{tokenizer.decode(input_ids)}')
+
+
+def core_process(encoded_docs, semaphore, seq_length, tokenizer, encoder, builder, output_idx_file):
+ """
+ core of Data Pack SFT processing
+ """
+ input_ids_key = 'input_ids'
+
+ proc_start = time.time()
+ total_bytes_processed = 0
+ pbar = tqdm.tqdm()
+ sentence_droped = 0
+ loss_token_cnt = 0
+
+ print("PRINT BEFORE STREAM PROCESS DATA")
+
+ print_example_count = 0
+ for i, (doc, bytes_processed) in enumerate(encoded_docs, start=1):
+ total_bytes_processed += bytes_processed
+
+ # release semaphore so `yield_from_files` can add another file to the buffer
+ semaphore.release()
+
+ # add each tokenized document / sentence,
+ # For sft, each document has only one sample
+ input_ids_sentence = doc[input_ids_key][0]
+ if len(input_ids_sentence) < 1:
+ sentence_droped += 1
+ continue
+
+ builder.add_item(np.array(input_ids_sentence, dtype=builder.dtype))
+ builder.end_document()
+ #builder.finalize_without_close(output_idx_file)
+ #builder.add_item_and_end_document_and_finalize(np.array(input_ids_sentence, dtype=builder.dtype), output_idx_file)
+
+ # print the first packed sample as example
+ if print_example_count < 1:
+ print_example_doc(input_ids_sentence, tokenizer)
+ print_example_count += 1
+
+ # log progress
+ if i % 100 == 0:
+ current = time.time()
+ elapsed = current - proc_start
+ mbs = total_bytes_processed / elapsed / 1024 / 1024
+ pbar.set_description(
+ f"Processed {i} documents ({i / elapsed} docs/s, {mbs} MB/s)."
+ )
+ if i != 0:
+ pbar.update(100)
+
+ # 尾部处理
+ builder.finalize(output_idx_file)
+
+ print(Fore.RED + "\ndroped docs: {}".format(sentence_droped))
+
+
+def process_dataset(dataset_path, output_path, model_path, parallel_num, seq_length, dataset_name, sample_percent):
+ """
+ Re-organize samples in the given data path into a Data Pack file.
+ """
+
+ # get all jsonl files and corresponding reading handler
+ files = glob.glob(os.path.join(dataset_path, '**/*.jsonl'), recursive=True)
+
+ # build a semaphore object to stop `yield_from_files` from getting ahead
+ # of encoder.encode and hence building up memory
+ semaphore = Semaphore(1000 + parallel_num)
+
+ # build sample iterator
+ sample_iterator = yield_from_files2(files, semaphore, sample_percent)
+
+ # load tokenizer
+ # tokenizer = load_tokenizer(model_path, tokenizer_type)
+ tokenizer = init_tokenizer(model_path)
+ print('TOKEN of id=2:', tokenizer.convert_ids_to_tokens(2))
+ print('ID of
@@ -21,7 +21,11 @@
-
-
-[[中文]](README_cn.md) [**English**]
-
-## 1. Updates
-
-🔥 MFTCoder supports QLoRA/LoRA using Huggingface accelerate + DeepSpeed Framework;
-
-🔥 MFTCoder supports Multiple Task Finetuning, which is able to balance diffenrent tasks in data level.
-
-🔥 MFTCoder supports finetuning multiple mainstream open-source base models: codellama, llama2, llama, starcoder, codegeex2, chatglm2, qwen.
-
-## 2. Data Format
-### 2.1 Training Data Format
-The training data is in a uniformed JSONL format, in which each line of data has the following JSON format. The "chat_rounds" field is required, and other fields can be added or removed based on specific needs.
-
-```json
-{
- "id":0,
- "data_name":"code-helper",
- "chat_rounds":[
- {
- "role": "system",
- "content": "You are a expert in coding and help answer code questions",
- "chat_round_id": 0
- },
- {
- "role": "human",
- "content": "Write a python function of quick sort",
- "chat_round_id": 1
- },
- {
- "role": "bot",
- "content": "Below is the function of quick sort: ...",
- "chat_round_id": 1
- },
- {
- "role": "human",
- "content": "Explain the code",
- "chat_round_id": 2
- },
- {
- "role": "bot",
- "content": "OK, this code ...",
- "chat_round_id": 2
- }
- ]
-}
-```
-
-### 2.2 Inference Data Format
-The inference data contains strings concatenated by conversation data(system, human and bot contents) in the training data format.
-It is used as the data "seen"(before tokenization) by the model in training process.
-It is used as input during the inference process as well.
-Here is an example format of the concatenated string:
-
-```python
-"""
-<|role_start|>system<|role_end|>System instruction
-<|role_start|>human<|role_end|>Human 1st round input
-<|role_start|>bot<|role_end|>Bot 1st round output
-<|role_start|>human<|role_end|>Human 2nd round input
-<|role_start|>bot<|role_end|>Bot 2nd round output
-...
-...
-...
-<|role_start|>human<|role_end|>Human nth round input
-<|role_start|>bot<|role_end|>{Bot output to be genreated}
-"""
-```
-When applying inference, you always make your input string end with "<|role_start|>bot<|role_end|>" to request the model generating answers.
-
-
-## 3. Model Training
-Currently, the "MFTCoder/mft_peft_hf" codebase supports QLoRA instruction fine-tuning, and LoRA instruction fine-tuning.
-In theory, this project can be used to train any publicly available model in the HuggingFace Format.
-
-Here are some excellent pre-trained models weights available on Huggingface that can be finetuned with this codebase:
-
-🤗 [Latest code pre-trained SOTA, CodeLlama-34b-Python](https://huggingface.co/codellama/CodeLlama-34b-Python-hf) : code-llama-34b, code-llama-34b-python, a new SOTA base model.
-
-🤗 [Best 10B level pre-trained Code LLM, Starcoder:](https://huggingface.co/bigcode/starcoder) wizardCoder-15B, PanGu-coder2, and other previous SOTA were trained on it.
-
-🤗 [Multilingual powerhouse, Qwen-7b](https://huggingface.co/Qwen/Qwen-7B): Suitable for multilingual tasks, including Chinese tasks, for instruction fine-tuning.
-
-You can find the implementations in the ```mft_peft_hf/src``` directory. The entry directory for fine-tuning training is ```mft_peft_hf/src/pefts```, and the entry file for training is ```mft_peft_hf/src/pefts/mft_accelerate.py```.
-Configurations are stored in the ```mft_peft_hf/src/pefts/configs``` directory for easy management and modification.
-
-### 3.1 Tokenization
-During training, we concatenate multi-turn dialogues into the following format (also known as the inference data format mentioned earlier) and then tokenize it. In this format, <|role_start|>human<|role_end|> represents the human input (i.e., prompt), <|role_start|>bot<|role_end|> represents the bot output, and represents the eos_token.
-You can modify and `````` replace the eos_token based on different models' requirements.
-
-Here is an example of the concatenated format with prompts:
-```
-"<|role_start|>human<|role_end|>input1<|role_start|>bot<|role_end|>target1<|role_start|>human<|role_end|>input2<|role_start|>bot<|role_end|>target2...
-```
-During the calculation of loss, we use a ```loss mask``` to ensure that the loss from the input part does not contribute to parameter updates. Only the loss from the ```target``` part is used for updating parameters.
-This approach takes full advantage of the benefits of model parallelism, making training more efficient. It also leverages the characteristic of decoder-only models with left-to-right attention.
-By including all target parts from multiple turns in a single training iteration, the training process becomes more efficient.
-
-
-### 3.2 LoRA/QLoRA
-You can refer to the Lora paper for details about LoRA：[LORA: LOW-RANK ADAPTATION OF LARGE LANGUAGE MODELS](https://arxiv.org/pdf/2106.09685.pdf)
-You can refer to the Qlora paper for details about QLoRA：[QLORA: Efficient Finetuning of Quantized LLMs](https://arxiv.org/pdf/2305.14314.pdf)
-
-QLoRA (Quantized LoRA) is a method that combines 4-bit nf4 quantization and additional adapters to achieve a balance between reducing GPU memory consumption and approaching the performance of full-parameter fine-tuning.
-
-According to the QLoRA paper, this method enables fine-tuning of a 33B model on a single V100 GPU while achieving performance close to that of full-parameter fine-tuning.
-
-To perform LoRA/QLoRA fine-tuning, you can execute the following command:
-```bash
-cd mft_peft_hf/src/pefts
-
-accelerate launch --config_file accelerate_ds_config.yaml mft_accelerate.py --train_config configs/starcoder_train_config.json
-```
-The main parameter explanations for the ```configs/*_train_config``` are as follows. You can modify these parameters according to your needs:
-
-- **load_raw_dataset**: Need to be true at present. Only JSONL format is supported.
-
-- **data_paths**: Input data paths in a String of list format, e.g., "[path1,path2,path3]". Each path represents a task directory and each task directory contains one or more JSONL data files.
-
-- **output_dir**: Training output directory to store checkpoints, Lora adapter, etc.
-
-- **tb_dir**: TensorBoard directory to store logs, metrics, etc.
-
-- **model_type**: Type of the model to train, e.g., "llama | starcoder | chatglm2 | qwen | gpt_neox".
-
-- **peft_type**: either "lora" or "qlora".
-
-- **lora_rank**: Rank value for Lora.
-
-- **lora_alpha**: Alpha value for Lora.
-
-- **lora_dropout**: Dropout rate for Lora.
-
-- **quantization**: Whether to use quantization."4bit" or "8bit", or null. For QLoRA, it is recommended to use 4-bit quantization.
-
-- **pretrained_model_path**: Local/Shared disk path or model name on HuggingFace for the pre-trained model.
-
-- **weighted_loss_mode**: Loss weighting method for multitask training. "case3" is recommended at present.
-
-- **padding_mode**: The way tokenized data is set. "padding" means padding for each sample to seq_length, "pack" means putting samples into seq_length as many as possible.
-
-- **num_train_epochs**: Number of training epochs.
-
-- **per_device_train_batch_size**: Batch size per GPU for training.
-
-- **per_device_eval_batch_size**: Batch size per GPU for evaluation.
-
-- **gradient_accumulation_steps**: Number of gradient accumulation steps. Global batch size is calculated as num_gpus * per_device_train_batch_size * gradient_accumulation_steps.
-
-- **learning_rate**: Initial Learning rate. For full-parameter fine-tuning, it is recommended to use a smaller value such as 1e-5 or 5e-6. For QLoRA, a larger learning rate is generally used, such as 1e-4 or 2e-4.
-
-- **min_lr**: Minimum learning rate. Usually set to one-tenth of the learning rate.
-
-- **seq_length**: Maximum input sequence length during training.
-
-- **log_interval**: Log training loss every ```log_interval``` steps.
-
-- **checkpointing_steps**: Save a checkpoint every ```checkpointing_steps``` steps.
-
-- **evaluation_steps**: Evaluate on the validation set every ```evaluation_steps``` steps.
-
-- **early_stopping**: Enable early stopping or not.
-
-- **early_stopping_stall_num**: Number of evaluation points without improvement which triggers early stopping.
-
-- **lr_scheduler_type**: Type of learning rate scheduler. "cosine" is a good choice already.
-
-- **num_warmup_steps**: Number of warm-up steps to gradually increase the learning rate.
-
-- **seed**: Random seed for reproducibility.
-
-
-## 4. Model Usage
-
-### 4.1 Merge Adaptor weights
-Using LoRA or QLoRA for training, this project only saves the weights and configuration files of the adapters.
-To merge the adapter weights with the base model, see ```src/pefts/merge_base_and_lora_to_hf.py```
-
-### 4.2 Inference demo
-Here is the script for inference on our trained models, which is compatible with most HuggingFace models:
-```python
-from transformers import (
- AutoTokenizer,
- AutoModelForCausalLM,
-)
-tokenizer = AutoTokenizer.from_pretrained(mode_name_or_path, trust_remote_code=True, use_fast=False, legacy=False)
-tokenizer.padding_side = "left"
-tokenizer.pad_token_id = tokenizer.convert_tokens_to_ids("