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.2 KiB69 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
9
10
11@ModelBase.register("GPTRefactForCausalLM")
12@ModelBase.example("smallcloudai/Refact-1_6-base")
13class RefactModel(TextModel):
14    model_arch = gguf.MODEL_ARCH.REFACT
15
16    def set_vocab(self):
17        super().set_vocab()
18
19        # TODO: how to determine special FIM tokens automatically?
20        special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=False,
21                                          special_token_types = ['prefix', 'suffix', 'middle', 'eot'])
22        special_vocab._set_special_token("prefix", 1)
23        special_vocab._set_special_token("suffix", 3)
24        special_vocab._set_special_token("middle", 2)
25        special_vocab.chat_template = None  # do not add it twice
26        special_vocab.add_to_gguf(self.gguf_writer)
27
28    def set_gguf_parameters(self):
29        hidden_dim = self.hparams["n_embd"]
30        inner_dim = 4 * hidden_dim
31        hidden_dim = int(2 * inner_dim / 3)
32        multiple_of = 256
33        ff_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
34
35        # refact uses Alibi. So this is from config.json which might be used by training.
36        self.gguf_writer.add_context_length(self.hparams["n_positions"])
37        self.gguf_writer.add_embedding_length(self.hparams["n_embd"])
38
39        self.gguf_writer.add_feed_forward_length(ff_dim)
40        self.gguf_writer.add_block_count(self.block_count)
41        self.gguf_writer.add_head_count(self.hparams["n_head"])
42        self.gguf_writer.add_head_count_kv(1)
43        self.gguf_writer.add_layer_norm_rms_eps(self.hparams["layer_norm_epsilon"])
44        self.gguf_writer.add_file_type(self.ftype)
45
46    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
47        hidden_dim = self.hparams["n_embd"]
48        inner_dim = 4 * hidden_dim
49        hidden_dim = int(2 * inner_dim / 3)
50        multiple_of = 256
51        ff_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
52        n_head = self.hparams["n_head"]
53        n_head_kv = 1
54        head_dim = self.hparams["n_embd"] // n_head
55
56        if bid is not None:
57            if name == f"transformer.h.{bid}.attn.kv.weight":
58                yield from super().modify_tensors(data_torch[:n_head_kv * head_dim], self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_K, bid), bid)
59                yield from super().modify_tensors(data_torch[n_head_kv * head_dim:], self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_V, bid), bid)
60                return
61            if name == f"transformer.h.{bid}.attn.q.weight":
62                yield from super().modify_tensors(data_torch, self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_Q, bid), bid)
63                return
64            if name == f"transformer.h.{bid}.mlp.gate_up_proj.weight":
65                yield from super().modify_tensors(data_torch[:ff_dim], self.format_tensor_name(gguf.MODEL_TENSOR.FFN_GATE, bid), bid)
66                yield from super().modify_tensors(data_torch[ff_dim:], self.format_tensor_name(gguf.MODEL_TENSOR.FFN_UP, bid), bid)
67                return
68
69        yield from super().modify_tensors(data_torch, name, bid)