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("StableLmForCausalLM", "StableLMEpochForCausalLM", "LlavaStableLMEpochForCausalLM")14class StableLMModel(TextModel):15    model_arch = gguf.MODEL_ARCH.STABLELM16 17    def set_vocab(self):18        if (self.dir_model / "tokenizer.json").is_file():19            self._set_vocab_gpt2()20        else:21            # StableLM 2 1.6B used to have a vocab in a similar format to Qwen's vocab22            self._set_vocab_qwen()23 24    def set_gguf_parameters(self):25        hparams = self.hparams26 27        self.gguf_writer.add_context_length(hparams["max_position_embeddings"])28        self.gguf_writer.add_embedding_length(hparams["hidden_size"])29        self.gguf_writer.add_block_count(self.block_count)30        self.gguf_writer.add_feed_forward_length(hparams["intermediate_size"])31        rotary_factor = self.rope_parameters["partial_rotary_factor"]32        self.gguf_writer.add_rope_dimension_count(int(rotary_factor * (hparams["hidden_size"] // hparams["num_attention_heads"])))33        self.gguf_writer.add_head_count(hparams["num_attention_heads"])34        self.gguf_writer.add_head_count_kv(hparams["num_key_value_heads"])35        self.gguf_writer.add_parallel_residual(hparams["use_parallel_residual"] if "use_parallel_residual" in hparams else True)36        self.gguf_writer.add_layer_norm_eps(self.find_hparam(["layer_norm_eps", "norm_eps"]))37        self.gguf_writer.add_file_type(self.ftype)38 39    _q_norms: list[dict[str, Tensor]] | None = None40    _k_norms: list[dict[str, Tensor]] | None = None41 42    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:43        n_head = self.hparams["num_attention_heads"]44        n_kv_head = self.hparams["num_key_value_heads"]45 46        if name.find("q_layernorm.norms") != -1:47            assert bid is not None48 49            if self._q_norms is None:50                self._q_norms = [{} for _ in range(self.block_count)]51 52            self._q_norms[bid][name] = data_torch53 54            if len(self._q_norms[bid]) >= n_head:55                return self._stack_qk_norm(bid, n_head, self._q_norms[bid], "q_layernorm")56            else:57                return58 59        if name.find("k_layernorm.norms") != -1:60            assert bid is not None61 62            if self._k_norms is None:63                self._k_norms = [{} for _ in range(self.block_count)]64 65            self._k_norms[bid][name] = data_torch66 67            if len(self._k_norms[bid]) >= n_kv_head:68                return self._stack_qk_norm(bid, n_kv_head, self._k_norms[bid], "k_layernorm")69            else:70                return71 72        yield from super().modify_tensors(data_torch, name, bid)73 74    def _stack_qk_norm(self, bid: int, n_head: int, norms: dict[str, Tensor], layer_name: str = "q_layernorm"):75        datas: list[Tensor] = []76        # extract the norms in order77        for xid in range(n_head):78            ename = f"model.layers.{bid}.self_attn.{layer_name}.norms.{xid}.weight"79            datas.append(norms[ename])80            del norms[ename]81        data_torch = torch.stack(datas, dim=0)82 83        merged_name = f"model.layers.{bid}.self_attn.{layer_name}.weight"84 85        yield from super().modify_tensors(data_torch, merged_name, bid)86 87    def prepare_tensors(self):88        super().prepare_tensors()89 90        if self._q_norms is not None or self._k_norms is not None:91            # flatten two `list[dict[str, Tensor]]` into a single `list[str]`92            norms = (93                [k for d in self._q_norms for k in d.keys()] if self._q_norms is not None else []94            ) + (95                [k for d in self._k_norms for k in d.keys()] if self._k_norms is not None else []96            )97            if len(norms) > 0:98                raise ValueError(f"Unprocessed norms: {norms}")99