Felipe97/llama-cpp-compiled

1from __future__ import annotations2 3from typing import Iterable, TYPE_CHECKING4 5import torch6 7if TYPE_CHECKING:8    from torch import Tensor9 10from .base import ModelBase, TextModel, gguf, logger11 12 13@ModelBase.register("JambaForCausalLM")14@ModelBase.example("ai21labs/Jamba-v0.1")15class JambaModel(TextModel):16    model_arch = gguf.MODEL_ARCH.JAMBA17 18    def set_vocab(self):19        if (self.dir_model / "tokenizer.model").is_file():20            self._set_vocab_sentencepiece()21        else:22            self._set_vocab_llama_hf()23            self.gguf_writer.add_add_space_prefix(False)24 25    def set_gguf_parameters(self):26        d_model = self.find_hparam(["hidden_size", "mamba_d_model"])27        d_conv  = self.find_hparam(["mamba_d_conv"],  optional=True) or 428        d_inner = self.hparams["mamba_expand"] * d_model29        d_state = self.find_hparam(["mamba_d_state"], optional=True) or 1630        # ceiling division31        # ref: https://stackoverflow.com/a/17511341/2282786332        # ref: https://github.com/state-spaces/mamba/blob/ce59daea3a090d011d6476c6e5b97f6d58ddad8b/mamba_ssm/modules/mamba_simple.py#L5833        dt_rank      = self.find_hparam(["mamba_dt_rank"], optional=True) or -(d_model // -16)34        rms_norm_eps = self.find_hparam(["layer_norm_epsilon", "rms_norm_eps"], optional=True) or 1e-635        n_kv_head = self.hparams["num_key_value_heads"]36        attn_offset = self.hparams["attn_layer_offset"]37        attn_period = self.hparams["attn_layer_period"]38        n_kv_vec = [0 for _ in range(attn_offset)] + [39            n_kv_head if (i - attn_offset) % attn_period == 0 else 0 for i in range(attn_offset, self.block_count)40        ]41 42        self.gguf_writer.add_block_count(self.block_count)43        self.gguf_writer.add_context_length(self.find_hparam(["max_position_embeddings", "n_ctx"]))44        self.gguf_writer.add_embedding_length(d_model)45        self.gguf_writer.add_feed_forward_length(self.hparams["intermediate_size"])46        self.gguf_writer.add_head_count(self.hparams["num_attention_heads"])47        self.gguf_writer.add_head_count_kv(n_kv_vec)48        self.gguf_writer.add_ssm_conv_kernel(d_conv)49        self.gguf_writer.add_ssm_inner_size(d_inner)50        self.gguf_writer.add_ssm_state_size(d_state)51        self.gguf_writer.add_ssm_time_step_rank(dt_rank)52        self.gguf_writer.add_layer_norm_rms_eps(rms_norm_eps)53        self.gguf_writer.add_expert_count(self.find_hparam(["num_local_experts", "num_experts"]))54        self.gguf_writer.add_expert_used_count(self.find_hparam(["num_experts_per_tok", "num_experts_per_token"]))55        self.gguf_writer.add_file_type(self.ftype)56 57    _experts: list[dict[str, Tensor]] | None = None58 59    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:60 61        # Mini-Jamba62        name = name.replace(".moe.", ".feed_forward.")63        if bid is not None:64            moe_offset = self.hparams["expert_layer_offset"]65            moe_period = self.hparams["expert_layer_period"]66 67            if not (bid >= moe_offset and (bid - moe_offset) % moe_period == 0):68                name = name.replace(".experts.0.", ".")69 70        # process the experts separately71        if ".feed_forward.experts." in name:72            n_experts = self.find_hparam(["num_local_experts", "num_experts"])73 74            assert bid is not None75 76            if self._experts is None:77                self._experts = [{} for _ in range(self.block_count)]78 79            self._experts[bid][name] = data_torch80 81            if len(self._experts[bid]) >= n_experts * 3:82 83                # merge the experts into a single 3d tensor84                for wid in ["down_proj", "gate_proj", "up_proj"]:85                    datas: list[Tensor] = []86 87                    for xid in range(n_experts):88                        ename = f"model.layers.{bid}.feed_forward.experts.{xid}.{wid}.weight"89                        datas.append(self._experts[bid][ename])90                        del self._experts[bid][ename]91 92                    data_torch = torch.stack(datas, dim=0)93 94                    # using the same merged name as qwen2moe95                    merged_name = f"model.layers.{bid}.mlp.experts.{wid}.weight"96 97                    new_name = self.map_tensor_name(merged_name)98 99                    yield new_name, data_torch100            return101 102        new_name = self.map_tensor_name(name)103 104        if self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.SSM_CONV1D, bid):105            data_torch = data_torch.squeeze()106 107        if name.endswith(".A_log"):108            logger.debug("A_log --> A ==> " + new_name)109            data_torch = -torch.exp(data_torch)110 111        yield (new_name, data_torch)112 113    def prepare_tensors(self):114        super().prepare_tensors()115 116        if self._experts is not None:117            # flatten `list[dict[str, Tensor]]` into `list[str]`118            experts = [k for d in self._experts for k in d.keys()]119            if len(experts) > 0:120                raise ValueError(f"Unprocessed experts: {experts}")121