Felipe97/llama-cpp-compiled

1from __future__ import annotations2 3from typing import Iterable, TYPE_CHECKING4 5if TYPE_CHECKING:6    from torch import Tensor7 8from .base import ModelBase, TextModel, gguf, logger9 10 11@ModelBase.register("DbrxForCausalLM")12@ModelBase.example("alpindale/dbrx-instruct")13class DbrxModel(TextModel):14    model_arch = gguf.MODEL_ARCH.DBRX15 16    def set_gguf_parameters(self):17        ffn_config = self.hparams["ffn_config"]18        attn_config = self.hparams["attn_config"]19        self.gguf_writer.add_block_count(self.block_count)20 21        self.gguf_writer.add_context_length(self.hparams["max_seq_len"])22        self.gguf_writer.add_embedding_length(self.hparams["d_model"])23        self.gguf_writer.add_feed_forward_length(ffn_config["ffn_hidden_size"])24 25        self.gguf_writer.add_head_count(self.hparams["n_heads"])26        self.gguf_writer.add_head_count_kv(attn_config["kv_n_heads"])27 28        self.gguf_writer.add_rope_freq_base(attn_config["rope_theta"])29 30        self.gguf_writer.add_clamp_kqv(attn_config["clip_qkv"])31 32        self.gguf_writer.add_expert_count(ffn_config["moe_num_experts"])33        self.gguf_writer.add_expert_used_count(ffn_config["moe_top_k"])34 35        self.gguf_writer.add_layer_norm_eps(1e-5)36 37        self.gguf_writer.add_file_type(self.ftype)38        logger.info(f"gguf: file type = {self.ftype}")39 40    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:41        n_expert = self.hparams["ffn_config"]["moe_num_experts"]42        n_ff = self.hparams["ffn_config"]["ffn_hidden_size"]43        n_embd = self.hparams["d_model"]44 45        # Specific behavior for experts tensors: suffix .weight, view as 3D and transpose46        # original implementation expects (n_expert, n_ff, n_embd) for all experts weights47        # But llama.cpp moe graph works differently48        # AND the dimensions in ggml are typically in the reverse order of the pytorch dimensions49        # so (n_expert, n_ff, n_embd) in pytorch is {n_embd, n_ff, n_expert} in ggml_tensor50        exp_tensor_names = {"ffn.experts.mlp.w1": None,       # LLM_TENSOR_FFN_GATE_EXPS ggml_tensor->ne{n_embd, n_ff,   n_expert}51                            "ffn.experts.mlp.w2": (0, 2, 1),  # LLM_TENSOR_FFN_DOWN_EXPS ggml_tensor->ne{n_ff,   n_embd, n_expert}52                            "ffn.experts.mlp.v1": None}       # LLM_TENSOR_FFN_UP_EXPS   ggml_tensor->ne{n_embd, n_ff,   n_expert}53        experts = False54 55        for exp_tensor_name in exp_tensor_names.keys():56            if name.find(exp_tensor_name) != -1 and name.find(".weight") == -1:57                experts = True58                data_torch = data_torch.view(n_expert, n_ff, n_embd)59                if (permute_tensor := exp_tensor_names[exp_tensor_name]) is not None:60                    data_torch = data_torch.permute(*permute_tensor)61                break62 63        # map tensor names64        # In MoE models the ffn tensors are typically most of the model weights,65        # and need to be quantizable. Quantize expects tensor names to be suffixed by .weight.66        # Every other model has the weight names ending in .weight,67        # let's assume that is the convention which is not the case for dbrx:68        # https://huggingface.co/databricks/dbrx-instruct/blob/main/model.safetensors.index.json#L1569        new_name = self.map_tensor_name(name if not experts else name + ".weight", try_suffixes=(".weight",))70 71        yield from super().modify_tensors(data_torch, new_name, bid)72 73    def tensor_force_quant(self, name: str, new_name: str, bid: int | None, n_dims: int) -> gguf.GGMLQuantizationType | bool:74        del name, new_name, bid  # unused75 76        return n_dims > 177