Brunobkr/llama.cpp_AlgMor24_github

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")12class DbrxModel(TextModel):13    model_arch = gguf.MODEL_ARCH.DBRX14 15    def set_gguf_parameters(self):16        ffn_config = self.hparams["ffn_config"]17        attn_config = self.hparams["attn_config"]18        self.gguf_writer.add_block_count(self.block_count)19 20        self.gguf_writer.add_context_length(self.hparams["max_seq_len"])21        self.gguf_writer.add_embedding_length(self.hparams["d_model"])22        self.gguf_writer.add_feed_forward_length(ffn_config["ffn_hidden_size"])23 24        self.gguf_writer.add_head_count(self.hparams["n_heads"])25        self.gguf_writer.add_head_count_kv(attn_config["kv_n_heads"])26 27        self.gguf_writer.add_rope_freq_base(attn_config["rope_theta"])28 29        self.gguf_writer.add_clamp_kqv(attn_config["clip_qkv"])30 31        self.gguf_writer.add_expert_count(ffn_config["moe_num_experts"])32        self.gguf_writer.add_expert_used_count(ffn_config["moe_top_k"])33 34        self.gguf_writer.add_layer_norm_eps(1e-5)35 36        self.gguf_writer.add_file_type(self.ftype)37        logger.info(f"gguf: file type = {self.ftype}")38 39    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:40        n_expert = self.hparams["ffn_config"]["moe_num_experts"]41        n_ff = self.hparams["ffn_config"]["ffn_hidden_size"]42        n_embd = self.hparams["d_model"]43 44        # Specific behavior for experts tensors: suffix .weight, view as 3D and transpose45        # original implementation expects (n_expert, n_ff, n_embd) for all experts weights46        # But llama.cpp moe graph works differently47        # AND the dimensions in ggml are typically in the reverse order of the pytorch dimensions48        # so (n_expert, n_ff, n_embd) in pytorch is {n_embd, n_ff, n_expert} in ggml_tensor49        exp_tensor_names = {"ffn.experts.mlp.w1": None,       # LLM_TENSOR_FFN_GATE_EXPS ggml_tensor->ne{n_embd, n_ff,   n_expert}50                            "ffn.experts.mlp.w2": (0, 2, 1),  # LLM_TENSOR_FFN_DOWN_EXPS ggml_tensor->ne{n_ff,   n_embd, n_expert}51                            "ffn.experts.mlp.v1": None}       # LLM_TENSOR_FFN_UP_EXPS   ggml_tensor->ne{n_embd, n_ff,   n_expert}52        experts = False53 54        for exp_tensor_name in exp_tensor_names.keys():55            if name.find(exp_tensor_name) != -1 and name.find(".weight") == -1:56                experts = True57                data_torch = data_torch.view(n_expert, n_ff, n_embd)58                if (permute_tensor := exp_tensor_names[exp_tensor_name]) is not None:59                    data_torch = data_torch.permute(*permute_tensor)60                break61 62        # map tensor names63        # In MoE models the ffn tensors are typically most of the model weights,64        # and need to be quantizable. Quantize expects tensor names to be suffixed by .weight.65        # Every other model has the weight names ending in .weight,66        # let's assume that is the convention which is not the case for dbrx:67        # https://huggingface.co/databricks/dbrx-instruct/blob/main/model.safetensors.index.json#L1568        new_name = self.map_tensor_name(name if not experts else name + ".weight", try_suffixes=(".weight",))69 70        yield from super().modify_tensors(data_torch, new_name, bid)71 72    def tensor_force_quant(self, name: str, new_name: str, bid: int | None, n_dims: int) -> gguf.GGMLQuantizationType | bool:73        del name, new_name, bid  # unused74 75        return n_dims > 176