Felipe97/llama-cpp-compiled

1from __future__ import annotations2 3import math4 5from typing import Any, Iterable, TYPE_CHECKING6 7import torch8 9if TYPE_CHECKING:10    from torch import Tensor11 12from .base import ModelBase, TextModel, gguf13 14 15@ModelBase.register("DeciLMForCausalLM")16@ModelBase.example("nvidia/Llama-3_1-Nemotron-51B-Instruct", "Deci/DeciLM-7B")17class DeciModel(TextModel):18    model_arch = gguf.MODEL_ARCH.DECI19 20    @staticmethod21    def _ffn_mult_to_intermediate_size(ffn_mult: float, n_embd: int) -> int:22        # DeciLM-specific code23        intermediate_size = int(2 * ffn_mult * n_embd / 3)24        return DeciModel._find_multiple(intermediate_size, 256)25 26    @staticmethod27    def _find_multiple(n: int, k: int) -> int:28        # DeciLM-specific code29        if n % k == 0:30            return n31        return n + k - (n % k)32 33    def __init__(self, *args, **kwargs):34        super().__init__(*args, **kwargs)35 36        if "block_configs" in self.hparams: # Llama-3_1-Nemotron-51B37            _block_configs: list[dict[str,Any]] = self.hparams["block_configs"]38            assert self.block_count == len(_block_configs)39            self._num_kv_heads = list()40            self._num_heads = list()41            _ffn_multipliers = list()42            # ***linear attention layer***43            # if n_heads_in_group is None and replace_with_linear is True44            # then _num_kv_heads[il] is 0 and _num_heads[il] is num_attention_heads45            # ***attention-free layer***46            # if n_heads_in_group is None and replace_with_linear is False47            # then _num_kv_heads[il] is 0 and _num_heads[il] is 048            # ***normal attention-layer***49            # if n_heads_in_group is not None, then50            # _num_kv_heads[il] is num_attention_head // n_heads_in_group and51            # _num_heads[il] is num_attention_head52            # ***dummy layer*** for nemotron 253B53            # if n_heads_in_group is None and ffn_mult is None54            # then _num_kv_heads[il] is 0 and _num_heads[il] is 0 and _ffn_dims is 055            for il in range(len(_block_configs)):56                if _block_configs[il]["attention"]["n_heads_in_group"] is None:57                    if _block_configs[il]["attention"]["replace_with_linear"] is True:58                        self._num_kv_heads.append(0)59                        self._num_heads.append(self.hparams["num_attention_heads"])60                    else:61                        self._num_kv_heads.append(0)62                        self._num_heads.append(0)63                else:64                    self._num_kv_heads.append(self.hparams["num_attention_heads"] // _block_configs[il]["attention"]["n_heads_in_group"])65                    self._num_heads.append(self.hparams["num_attention_heads"])66                if _block_configs[il]["ffn"]["ffn_mult"] is None: # dummy layer67                    _ffn_multipliers.append(0.0)68                else:69                    _ffn_multipliers.append(_block_configs[il]["ffn"]["ffn_mult"])70            assert self.block_count == len(self._num_kv_heads)71            assert self.block_count == len(self._num_heads)72            assert self.block_count == len(_ffn_multipliers)73            assert isinstance(self._num_kv_heads, list) and isinstance(self._num_kv_heads[0], int)74            assert isinstance(self._num_heads, list) and isinstance(self._num_heads[0], int)75            assert isinstance(_ffn_multipliers, list) and isinstance(_ffn_multipliers[0], float)76            self._ffn_dims: list[int] = [77                DeciModel._ffn_mult_to_intermediate_size(multiplier, self.hparams["hidden_size"])78                for multiplier in _ffn_multipliers79            ]80 81    def set_vocab(self):82        # Please change tokenizer_config.json of Llama-3_1-Nemotron-51B's83        # eos_token from '|eot_id|' to '|end_of_text|'84        if self.hparams.get("vocab_size", 128256) == 128256:85            tokens, toktypes, tokpre = self.get_vocab_base()86            self.gguf_writer.add_tokenizer_model("gpt2")87            self.gguf_writer.add_tokenizer_pre(tokpre)88            self.gguf_writer.add_token_list(tokens)89            self.gguf_writer.add_token_types(toktypes)90 91            special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True)92            special_vocab.add_to_gguf(self.gguf_writer)93        else:94            # DeciLM-7B95            self._set_vocab_llama_hf()96 97    def set_gguf_parameters(self):98        if "block_configs" in self.hparams: # Llama-3_1-Nemotron-51B99            assert self.block_count == len(self._num_kv_heads)100            assert self.block_count == len(self._num_heads)101            assert self.block_count == len(self._ffn_dims)102            if (rope_theta := self.rope_parameters.get("rope_theta")) is not None:103                self.gguf_writer.add_rope_freq_base(rope_theta)104            self.gguf_writer.add_head_count_kv(self._num_kv_heads)105            self.gguf_writer.add_head_count(self._num_heads)106            self.gguf_writer.add_feed_forward_length(self._ffn_dims)107            self.gguf_writer.add_block_count(self.block_count)108            self.gguf_writer.add_context_length(self.hparams["max_position_embeddings"])109            self.gguf_writer.add_embedding_length(self.hparams["hidden_size"])110            self.gguf_writer.add_layer_norm_rms_eps(self.hparams["rms_norm_eps"])111            self.gguf_writer.add_key_length(self.hparams["hidden_size"] // self.hparams["num_attention_heads"])112            self.gguf_writer.add_value_length(self.hparams["hidden_size"] // self.hparams["num_attention_heads"])113            self.gguf_writer.add_file_type(self.ftype)114        else: # DeciLM-7B115            super().set_gguf_parameters()116            if "num_key_value_heads_per_layer" in self.hparams: # DeciLM-7B117                self._num_kv_heads: list[int] = self.hparams["num_key_value_heads_per_layer"]118                assert self.block_count == len(self._num_kv_heads)119                self.gguf_writer.add_head_count_kv(self._num_kv_heads)120        hparams = self.hparams121        self.gguf_writer.add_vocab_size(hparams["vocab_size"])122 123        if (rope_dim := hparams.get("head_dim")) is None:124            rope_dim = hparams["hidden_size"] // hparams["num_attention_heads"]125        self.gguf_writer.add_rope_dimension_count(rope_dim)126 127    @staticmethod128    def permute(weights: Tensor, n_head: int, n_head_kv: int | None):129        if n_head_kv is not None and n_head != n_head_kv:130            n_head = n_head_kv131        return (weights.reshape(n_head, 2, weights.shape[0] // n_head // 2, *weights.shape[1:])132                .swapaxes(1, 2)133                .reshape(weights.shape))134 135    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:136        n_head = self.hparams["num_attention_heads"]137        if bid is not None:138            if "num_key_value_heads_per_layer" in self.hparams:139                n_kv_head = self.hparams["num_key_value_heads_per_layer"][bid]140            elif "block_configs" in self.hparams:141                n_kv_head = self._num_kv_heads[bid]142                n_head = self._num_heads[bid]143            else:144                n_kv_head = self.hparams.get("num_key_value_heads")145        else:146            n_kv_head = self.hparams.get("num_key_value_heads")147 148        if name.endswith(("q_proj.weight", "q_proj.bias")):149            data_torch = DeciModel.permute(data_torch, n_head, n_head)150        if name.endswith(("k_proj.weight", "k_proj.bias")):151            data_torch = DeciModel.permute(data_torch, n_head, n_kv_head)152        yield from super().modify_tensors(data_torch, name, bid)153 154    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:155        if rope_params := self.rope_parameters.get("full_attention", self.rope_parameters):156            if rope_params.get("rope_type", '').lower() == "llama3":157                base = rope_params.get("rope_theta", 10000.0)158                if (dim := self.hparams.get("head_dim")) is None:159                    dim = self.hparams["hidden_size"] // self.hparams["num_attention_heads"]160                freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))161 162                factor = rope_params.get("factor", 8.0)163                low_freq_factor = rope_params.get("low_freq_factor", 1.0)164                high_freq_factor = rope_params.get("high_freq_factor", 4.0)165                old_context_len = rope_params.get("original_max_position_embeddings", 8192)166 167                low_freq_wavelen = old_context_len / low_freq_factor168                high_freq_wavelen = old_context_len / high_freq_factor169                assert low_freq_wavelen != high_freq_wavelen170 171                rope_factors = []172                for freq in freqs:173                    wavelen = 2 * math.pi / freq174                    if wavelen < high_freq_wavelen:175                        rope_factors.append(1)176                    elif wavelen > low_freq_wavelen:177                        rope_factors.append(factor)178                    else:179                        smooth = (old_context_len / wavelen - low_freq_factor) / (high_freq_factor - low_freq_factor)180                        rope_factors.append(1 / ((1 - smooth) / factor + smooth))181 182                yield (self.format_tensor_name(gguf.MODEL_TENSOR.ROPE_FREQS), torch.tensor(rope_factors, dtype=torch.float32))183 184    def prepare_tensors(self):185        super().prepare_tensors()186