Brunobkr/llama.cpp_AlgMor24_github

1from __future__ import annotations2 3from typing import Any, Iterable, TYPE_CHECKING4 5if TYPE_CHECKING:6    from torch import Tensor7 8from .base import ModelBase, TextModel, gguf9 10 11@ModelBase.register("OpenELMForCausalLM")12class OpenELMModel(TextModel):13    model_arch = gguf.MODEL_ARCH.OPENELM14 15    @staticmethod16    def _make_divisible(v: float | int, divisor: int) -> int:17        # ref: https://huggingface.co/apple/OpenELM-270M-Instruct/blob/eb111ff2e6724348e5b905984063d4064d4bc579/configuration_openelm.py#L34-L3818        new_v = max(divisor, int(v + divisor / 2) // divisor * divisor)19        # Make sure that round down does not go down by more than 10%.20        if new_v < 0.9 * v:21            new_v += divisor22        return new_v23 24    def __init__(self, *args, **kwargs):25        super().__init__(*args, **kwargs)26 27        ffn_multipliers: list[float] = self.hparams["ffn_multipliers"]28        ffn_dim_divisor: int = self.hparams["ffn_dim_divisor"]29        self._n_embd: int = self.hparams["model_dim"]30        self._num_kv_heads: list[int] = self.hparams["num_kv_heads"]31        self._num_query_heads: list[int] = self.hparams["num_query_heads"]32        self._ffn_dims: list[int] = [33            OpenELMModel._make_divisible(multiplier * self._n_embd, ffn_dim_divisor)34            for multiplier in ffn_multipliers35        ]36        assert isinstance(self._num_kv_heads, list) and isinstance(self._num_kv_heads[0], int)37        assert isinstance(self._num_query_heads, list) and isinstance(self._num_query_heads[0], int)38 39    # Uses the tokenizer from meta-llama/Llama-2-7b-hf40    def set_vocab(self):41        try:42            self._set_vocab_sentencepiece()43        except FileNotFoundError:44            self._set_vocab_builtin("llama-spm", self.hparams["vocab_size"])45 46    def set_gguf_parameters(self):47        n_embd = self._n_embd48        head_dim = self.hparams["head_dim"]49        rot_pct = 1.050        assert self.block_count == len(self._num_kv_heads)51        assert self.block_count == len(self._num_query_heads)52        assert self.block_count == len(self._ffn_dims)53 54        self.gguf_writer.add_block_count(self.block_count)55        self.gguf_writer.add_context_length(self.hparams["max_context_length"])56        self.gguf_writer.add_embedding_length(n_embd)57        self.gguf_writer.add_feed_forward_length(self._ffn_dims)58        self.gguf_writer.add_head_count(self._num_query_heads)59        self.gguf_writer.add_head_count_kv(self._num_kv_heads)60        self.gguf_writer.add_rope_freq_base(self.hparams["rope_freq_constant"])61        # https://huggingface.co/apple/OpenELM-270M-Instruct/blob/c401df2/modeling_openelm.py#L3062        self.gguf_writer.add_layer_norm_rms_eps(1e-6)63        self.gguf_writer.add_rope_dimension_count(int(rot_pct * head_dim))64        self.gguf_writer.add_key_length(head_dim)65        self.gguf_writer.add_value_length(head_dim)66        self.gguf_writer.add_file_type(self.ftype)67 68    def find_hparam(self, keys: Iterable[str], optional: bool = False) -> Any:69        if "n_layers" in keys:70            return self.hparams["num_transformer_layers"]71 72        return super().find_hparam(keys, optional)73 74    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:75 76        # split ff77        if bid is not None and name == f"transformer.layers.{bid}.ffn.proj_1.weight":78            ff_dim = self._ffn_dims[bid]79            yield (self.format_tensor_name(gguf.MODEL_TENSOR.FFN_GATE, bid), data_torch[:ff_dim])80            yield (self.format_tensor_name(gguf.MODEL_TENSOR.FFN_UP, bid), data_torch[ff_dim:])81            return82 83        yield (self.map_tensor_name(name), data_torch)84