Brunobkr/llama.cpp_AlgMor24_github

1from __future__ import annotations2 3import json4 5from pathlib import Path6from typing import Callable, Iterable, TYPE_CHECKING7 8import torch9 10if TYPE_CHECKING:11    from torch import Tensor12 13from .base import ModelBase, TextModel, gguf, logger14 15 16@ModelBase.register("MambaForCausalLM", "MambaLMHeadModel", "FalconMambaForCausalLM")17class MambaModel(TextModel):18    model_arch = gguf.MODEL_ARCH.MAMBA19 20    def __init__(self, dir_model: Path, *args, **kwargs):21        # Avoid using AutoConfig for hparams22        hparams = kwargs.pop("hparams", None)23        if hparams is None:24            with open(dir_model / "config.json", "r", encoding="utf-8") as f:25                hparams = json.load(f)26        super().__init__(dir_model, *args, hparams=hparams, **kwargs)27 28    def set_vocab(self):29        vocab_size = self.hparams["vocab_size"]30        # Round vocab size to next multiple of 831        pad_vocab = self.hparams.get("pad_vocab_size_multiple", 8)32        # pad using ceiling division33        # ref: https://stackoverflow.com/a/17511341/2282786334        vocab_size = -(vocab_size // -pad_vocab) * pad_vocab35        self.hparams["vocab_size"] = vocab_size36 37        if (self.dir_model / "tokenizer.json").is_file():38            self._set_vocab_gpt2()39        elif (self.dir_model / "tokenizer.model").is_file():40            self._set_vocab_sentencepiece()41        else:42            # Use the GPT-NeoX tokenizer when no tokenizer files are present43            self._set_vocab_builtin("gpt-neox", vocab_size)44 45    def set_gguf_parameters(self):46        d_model = self.find_hparam(["hidden_size",       "d_model"])47        d_conv  = self.find_hparam(["conv_kernel",       "d_conv"],  optional=True) or 448        d_inner = self.find_hparam(["intermediate_size", "d_inner"], optional=True) or 2 * d_model49        d_state = self.find_hparam(["state_size",        "d_state"], optional=True) or 1650        # ceiling division51        # ref: https://stackoverflow.com/a/17511341/2282786352        # ref: https://github.com/state-spaces/mamba/blob/ce59daea3a090d011d6476c6e5b97f6d58ddad8b/mamba_ssm/modules/mamba_simple.py#L5853        dt_rank      = self.find_hparam(["time_step_rank",     "dt_rank"],      optional=True) or -(d_model // -16)54        rms_norm_eps = self.find_hparam(["layer_norm_epsilon", "rms_norm_eps"], optional=True) or 1e-555        use_dt_b_c_norm = False56        # For falconmamba we do apply RMS norm on B / DT and C layers57        if self.find_hparam(["model_type"], optional=True) in ("falcon_mamba",):58            use_dt_b_c_norm = True59        # Fail early for models which don't have a block expansion factor of 260        assert d_inner == 2 * d_model61 62        self.gguf_writer.add_context_length(2**20) # arbitrary value; for those who use the default63        self.gguf_writer.add_embedding_length(d_model)64        self.gguf_writer.add_feed_forward_length(0) # unused, but seemingly required when loading65        self.gguf_writer.add_head_count(0) # unused, but seemingly required when loading66        self.gguf_writer.add_block_count(self.block_count)67        self.gguf_writer.add_ssm_conv_kernel(d_conv)68        self.gguf_writer.add_ssm_inner_size(d_inner)69        self.gguf_writer.add_ssm_state_size(d_state)70        self.gguf_writer.add_ssm_time_step_rank(dt_rank)71        self.gguf_writer.add_layer_norm_rms_eps(rms_norm_eps)72        self.gguf_writer.add_ssm_dt_b_c_rms(use_dt_b_c_norm) # For classic Mamba we don't apply rms norm on B / DT layers73        self.gguf_writer.add_file_type(self.ftype)74 75    _tok_embd = None76 77    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:78        output_name = self.format_tensor_name(gguf.MODEL_TENSOR.OUTPUT)79        tok_embd_name = self.format_tensor_name(gguf.MODEL_TENSOR.TOKEN_EMBD)80 81        new_name = self.map_tensor_name(name)82 83        if name.endswith(".A_log"):84            logger.debug("A_log --> A ==> " + new_name)85            data_torch = -torch.exp(data_torch)86 87        # [4 1 8192 1] -> [4 8192 1 1]88        if self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.SSM_CONV1D, bid):89            data_torch = data_torch.squeeze()90 91        # assuming token_embd.weight is seen before output.weight92        if self._tok_embd is not None and new_name == output_name:93            if torch.equal(self._tok_embd, data_torch):94                logger.debug(f"{output_name} is equivalent to {tok_embd_name}, omitting")95                return96        elif new_name == tok_embd_name:97            self._tok_embd = data_torch98 99        yield from super().modify_tensors(data_torch, new_name, bid)100 101 102@ModelBase.register("Mamba2ForCausalLM")103class Mamba2Model(TextModel):104    model_arch = gguf.MODEL_ARCH.MAMBA2105 106    def __init__(self, dir_model: Path, *args, **kwargs):107        # Avoid using AutoConfig for hparams108        # It wrongly assumes all Mamba2 models are Mamba-Codestral-7B-v0.1109        hparams = kwargs.pop("hparams", None)110        if hparams is None:111            with open(dir_model / "config.json", "r", encoding="utf-8") as f:112                hparams = json.load(f)113        if "llm_config" in hparams:114            hparams["text_config"] = hparams["llm_config"]115        super().__init__(dir_model, *args, hparams=hparams, **kwargs)116        self.d_model = self.find_hparam(["hidden_size", "d_model", "dim"])117        self.expand = self.find_hparam(["mamba_expand", "expand"], optional=True) or 2118        self.d_inner = self.find_hparam(["mamba_d_ssm", "intermediate_size", "d_inner"], optional=True) or self.expand * self.d_model119        self.n_group = self.find_hparam(["n_groups"], optional=True) or 1120 121    def set_vocab(self):122        vocab_size = self.hparams["vocab_size"]123        # Round vocab size to next multiple of 16124        pad_vocab = self.hparams.get("pad_vocab_size_multiple", 16)125        # pad using ceiling division126        # ref: https://stackoverflow.com/a/17511341/22827863127        vocab_size = -(vocab_size // -pad_vocab) * pad_vocab128        self.hparams["vocab_size"] = vocab_size129 130        if (self.dir_model / "tokenizer.model").is_file():131            self._set_vocab_sentencepiece()132        elif (self.dir_model / "tokenizer.model.v3").is_file():133            # mamba-codestral134            raise NotImplementedError(f"Please rename {self.dir_model / 'tokenizer.model.v3'} to {self.dir_model / 'tokenizer.model'}")135        elif (self.dir_model / "tokenizer.json").is_file():136            self._set_vocab_gpt2()137        else:138            # Use the GPT-NeoX tokenizer when no tokenizer files are present139            self._set_vocab_builtin("gpt-neox", vocab_size)140 141    def set_gguf_parameters(self):142        d_conv  = self.find_hparam(["conv_kernel", "d_conv"],     optional=True) or 4143        d_state = self.find_hparam(["state_size",  "d_state"],    optional=True) or 128144        head_dim = self.find_hparam(["mamba_d_head", "head_dim"], optional=True) or 64145 146        rms_norm_eps = self.find_hparam(["layer_norm_epsilon", "rms_norm_eps"], optional=True) or 1e-5147 148        # skip the assertion for FalconH1 Model149        if self.model_arch != gguf.MODEL_ARCH.FALCON_H1:150            assert self.d_inner == self.expand * self.d_model151            assert self.d_inner % head_dim == 0152 153        self.gguf_writer.add_context_length(2**20)  # arbitrary value; for those who use the default154        self.gguf_writer.add_embedding_length(self.d_model)155        self.gguf_writer.add_feed_forward_length(0)  # unused, but seemingly required when loading156        self.gguf_writer.add_head_count(0)  # unused, but seemingly required when loading157        self.gguf_writer.add_block_count(self.block_count)158        self.gguf_writer.add_ssm_conv_kernel(d_conv)159        self.gguf_writer.add_ssm_inner_size(self.d_inner)160        self.gguf_writer.add_ssm_state_size(d_state)161        self.gguf_writer.add_ssm_time_step_rank(self.d_inner // head_dim)162        self.gguf_writer.add_ssm_group_count(self.n_group)163        self.gguf_writer.add_layer_norm_rms_eps(rms_norm_eps)164        self.gguf_writer.add_file_type(self.ftype)165 166    @classmethod167    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:168        name, gen = item169 170        if name.startswith(("model.backbone", "model.lm_head")):171            # map Mamba-Codestral-7B-v0.1 tensor names to the names used by Mamba-2172            name = name.removeprefix("model.")173 174        if name.endswith(".dt_bias"):175            name = name.rpartition(".dt_bias")[0] + ".dt_proj.bias"176 177        return super().filter_tensors((name, gen))178 179    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:180        new_name = self.map_tensor_name(name)181 182        if self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.SSM_CONV1D, bid):183            data_torch = data_torch.squeeze()184        elif any(self.match_model_tensor_name(new_name, t, bid, suffix="") for t in [185            gguf.MODEL_TENSOR.SSM_A,186            gguf.MODEL_TENSOR.SSM_D,187        ]):188            # unsqueeze A to use similar shape semantics as Mamba-1189            # (D is also unsqueezed, but for more straightforward broadcast internally)190            data_torch = data_torch.reshape((*data_torch.shape, 1))191        elif self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.SSM_NORM, bid):192            data_torch = data_torch.reshape((self.n_group, self.d_inner // self.n_group))193 194        if name.endswith(".A_log"):195            logger.debug("A_log --> A ==> " + new_name)196            data_torch = -torch.exp(data_torch)197 198        yield (new_name, data_torch)199