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 →
4.1 KiB107 linesraw
1from __future__ import annotations
2
3import math
4
5from typing import Callable, Iterable, TYPE_CHECKING
6
7if TYPE_CHECKING:
8    from torch import Tensor
9
10from .base import ModelBase, TextModel, gguf
11
12
13@ModelBase.register("Jais2ForCausalLM")
14# [TAG_HF_EXAMPLE_GATED] inceptionai/Jais-2-8B-Chat is gated
15# [TAG_HF_EXAMPLE_MISSING]
16class Jais2Model(TextModel):
17    model_arch = gguf.MODEL_ARCH.JAIS2
18
19    def set_gguf_parameters(self):
20        super().set_gguf_parameters()
21        hparams = self.hparams
22        head_dim = hparams.get("head_dim", hparams["hidden_size"] // hparams["num_attention_heads"])
23        self.gguf_writer.add_rope_dimension_count(head_dim)
24
25
26@ModelBase.register("JAISLMHeadModel")
27@ModelBase.example("inceptionai/jais-family-590m")
28class JaisModel(TextModel):
29    model_arch = gguf.MODEL_ARCH.JAIS
30
31    def __init__(self, *args, **kwargs):
32        super().__init__(*args, **kwargs)
33
34        # SwigLU activation
35        assert self.hparams["activation_function"] == "swiglu"
36        # ALiBi position embedding
37        assert self.hparams["position_embedding_type"] == "alibi"
38
39        # Embeddings scale
40        self.embeddings_scale = 1.0
41        if 'mup_embeddings_scale' in self.hparams:
42            self.embeddings_scale = self.hparams['mup_embeddings_scale']
43        elif 'embeddings_scale' in self.hparams:
44            self.embeddings_scale = self.hparams['embeddings_scale']
45        else:
46            assert False
47
48        self.width_scale = 1.0
49        if 'mup_output_alpha' in self.hparams:
50            assert 'mup_width_scale' in self.hparams
51            self.width_scale = self.hparams['mup_output_alpha'] * self.hparams['mup_width_scale']
52        elif 'width_scale' in self.hparams:
53            self.width_scale = self.hparams['width_scale']
54        else:
55            assert False
56
57        self.max_alibi_bias = 8.0
58
59    def set_vocab(self):
60        self._set_vocab_gpt2()
61
62    def set_gguf_parameters(self):
63        self.gguf_writer.add_block_count(self.block_count)
64        self.gguf_writer.add_context_length(self.hparams["n_positions"])
65        self.gguf_writer.add_embedding_length(self.hparams["n_embd"])
66        self.gguf_writer.add_feed_forward_length(self.hparams["n_inner"])
67        self.gguf_writer.add_head_count(self.hparams["n_head"])
68        self.gguf_writer.add_layer_norm_eps(self.hparams["layer_norm_epsilon"])
69        self.gguf_writer.add_file_type(self.ftype)
70
71    @classmethod
72    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
73        name, gen = item
74
75        # we don't need these
76        if name.endswith((".attn.bias")):
77            return None
78
79        return super().filter_tensors(item)
80
81    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
82        if name.endswith(("relative_pe.slopes")):
83            # Calculate max ALiBi bias (this is the inverse of the ALiBi calculation)
84            # Some other models has max_alibi_bias spelled out explicitly in the hyperparams,
85            # but Jais's PyTorch model simply precalculates the slope values and places them
86            # in relative_pes.slopes
87            n_head_closest_log2 = 2 ** math.floor(math.log2(self.hparams["n_head"]))
88            first_val = float(data_torch[0].item())
89            self.max_alibi_bias = -round(math.log2(first_val) * n_head_closest_log2)
90
91            return
92
93        if name.endswith((".c_attn.weight", ".c_proj.weight", ".c_fc.weight", ".c_fc2.weight")):
94            data_torch = data_torch.transpose(1, 0)
95
96        new_name = self.map_tensor_name(name)
97
98        if new_name == self.format_tensor_name(gguf.MODEL_TENSOR.TOKEN_EMBD):
99            yield from super().modify_tensors(data_torch * self.embeddings_scale, new_name, bid)
100        elif new_name == self.format_tensor_name(gguf.MODEL_TENSOR.OUTPUT):
101            yield from super().modify_tensors(data_torch * self.width_scale, new_name, bid)
102        else:
103            yield from super().modify_tensors(data_torch, new_name, bid)
104
105    def prepare_tensors(self):
106        super().prepare_tensors()
107        self.gguf_writer.add_max_alibi_bias(self.max_alibi_bias)