char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Xuan-Son Nguyenconvert: add @ModelBase.example (#27208)fa88ae936

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
3.6 KiB76 linesraw
1from __future__ import annotations
2
3from typing import Iterable, TYPE_CHECKING
4
5if TYPE_CHECKING:
6    from torch import Tensor
7
8from .base import ModelBase, TextModel, gguf, logger
9
10
11@ModelBase.register("DbrxForCausalLM")
12@ModelBase.example("alpindale/dbrx-instruct")
13class DbrxModel(TextModel):
14    model_arch = gguf.MODEL_ARCH.DBRX
15
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 transpose
46        # original implementation expects (n_expert, n_ff, n_embd) for all experts weights
47        # But llama.cpp moe graph works differently
48        # AND the dimensions in ggml are typically in the reverse order of the pytorch dimensions
49        # so (n_expert, n_ff, n_embd) in pytorch is {n_embd, n_ff, n_expert} in ggml_tensor
50        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 = False
54
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 = True
58                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                break
62
63        # map tensor names
64        # 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#L15
69        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  # unused
75
76        return n_dims > 1