Brunobkr/llama.cpp_AlgMor24_github

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, gguf11 12 13@ModelBase.register("FalconForCausalLM", "RWForCausalLM")14class FalconModel(TextModel):15    model_arch = gguf.MODEL_ARCH.FALCON16 17    def set_gguf_parameters(self):18        n_head = self.hparams.get("num_attention_heads")19        if n_head is None:20            n_head = self.hparams["n_head"]  # old name21 22        n_head_kv = self.hparams.get("num_kv_heads")23        if n_head_kv is None:24            n_head_kv = self.hparams.get("n_head_kv", 1)  # old name25 26        self.gguf_writer.add_context_length(2048)  # not in config.json27        self.gguf_writer.add_tensor_data_layout("jploski")  # qkv tensor transform28        self.gguf_writer.add_embedding_length(self.hparams["hidden_size"])29        self.gguf_writer.add_feed_forward_length(4 * self.hparams["hidden_size"])30        self.gguf_writer.add_block_count(self.block_count)31        self.gguf_writer.add_head_count(n_head)32        self.gguf_writer.add_head_count_kv(n_head_kv)33        self.gguf_writer.add_layer_norm_eps(self.hparams["layer_norm_epsilon"])34        self.gguf_writer.add_file_type(self.ftype)35 36    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:37        # QKV tensor transform38        # The original query_key_value tensor contains n_head_kv "kv groups",39        # each consisting of n_head/n_head_kv query weights followed by one key40        # and one value weight (shared by all query heads in the kv group).41        # This layout makes it a big pain to work with in GGML.42        # So we rearrange them here,, so that we have n_head query weights43        # followed by n_head_kv key weights followed by n_head_kv value weights,44        # in contiguous fashion.45        # ref: https://github.com/jploski/ggml/blob/falcon40b/examples/falcon/convert-hf-to-ggml.py46 47        if "query_key_value" in name:48            n_head = self.find_hparam(["num_attention_heads", "n_head"])49            n_head_kv = self.find_hparam(["num_kv_heads", "n_head_kv"], optional=True) or 150            head_dim = self.hparams["hidden_size"] // n_head51 52            qkv = data_torch.view(n_head_kv, n_head // n_head_kv + 2, head_dim, head_dim * n_head)53            q = qkv[:, :-2].reshape(n_head * head_dim, head_dim * n_head)54            k = qkv[:, [-2]].reshape(n_head_kv * head_dim, head_dim * n_head)55            v = qkv[:, [-1]].reshape(n_head_kv * head_dim, head_dim * n_head)56            data_torch = torch.cat((q, k, v)).reshape_as(data_torch)57 58        yield from super().modify_tensors(data_torch, name, bid)59