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