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 →
14.1 KiB306 linesraw
1from __future__ import annotations
2
3from typing import Callable, 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("Rwkv6ForCausalLM")
14@ModelBase.example("RWKV/v6-Finch-1B6-HF")
15class Rwkv6Model(TextModel):
16    model_arch = gguf.MODEL_ARCH.RWKV6
17
18    def set_vocab(self):
19        self._set_vocab_rwkv_world()
20
21    def set_gguf_parameters(self):
22        head_size = self.hparams["head_size"]
23        hidden_size = self.hparams["hidden_size"]
24        layer_norm_eps = self.hparams["layer_norm_epsilon"]
25        rescale_every_n_layers = self.hparams["rescale_every"]
26        intermediate_size = self.hparams["intermediate_size"] if self.hparams["intermediate_size"] is not None else int((hidden_size * 3.5) // 32 * 32)
27        time_mix_extra_dim = 64 if hidden_size == 4096 else 32
28        time_decay_extra_dim = 128 if hidden_size == 4096 else 64
29
30        # RWKV isn't context limited
31        self.gguf_writer.add_context_length(1048576)
32        self.gguf_writer.add_embedding_length(hidden_size)
33        self.gguf_writer.add_block_count(self.block_count)
34        self.gguf_writer.add_layer_norm_eps(layer_norm_eps)
35        self.gguf_writer.add_rescale_every_n_layers(rescale_every_n_layers)
36        self.gguf_writer.add_wkv_head_size(head_size)
37        self.gguf_writer.add_time_mix_extra_dim(time_mix_extra_dim)
38        self.gguf_writer.add_time_decay_extra_dim(time_decay_extra_dim)
39        self.gguf_writer.add_feed_forward_length(intermediate_size)
40        self.gguf_writer.add_file_type(self.ftype)
41
42        # required by llama.cpp, unused
43        self.gguf_writer.add_head_count(0)
44
45    lerp_weights: dict[int, dict[str, Tensor]] = {}
46
47    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
48        new_name = self.map_tensor_name(name)
49
50        if not (new_name.endswith(".weight") or new_name.endswith(".bias")):
51            new_name += ".weight"
52
53        if new_name.endswith("time_mix_w1.weight") or new_name.endswith("time_mix_decay_w1.weight") or new_name.endswith("time_mix_decay_w2.weight"):
54            data_torch = data_torch.transpose(0, 1)
55
56        if new_name.endswith("time_mix_w2.weight"):
57            data_torch = data_torch.permute(0, 2, 1)
58
59        if new_name.endswith("time_mix_decay.weight") or "lerp" in new_name:
60            data_torch = data_torch.squeeze()
61
62        try:
63            rescale_every_n_layers = self.hparams["rescale_every"]
64            if rescale_every_n_layers > 0:
65                if new_name.endswith("time_mix_output.weight") or new_name.endswith("channel_mix_value.weight"):
66                    data_torch = data_torch.div_(2 ** int(bid // rescale_every_n_layers))
67        except KeyError:
68            pass
69
70        # concat time_mix_lerp weights to reduce some cpu overhead
71        # also reduces the number of tensors in the model
72        if bid is not None and "time_mix_lerp" in new_name and "time_mix_lerp_x" not in new_name:
73            try:
74                self.lerp_weights[bid][new_name] = data_torch
75            except KeyError:
76                self.lerp_weights[bid] = {new_name: data_torch}
77            if all(f"blk.{bid}.time_mix_lerp_{i}.weight" in self.lerp_weights[bid].keys() for i in ["w", "k", "v", "r", "g"]):
78                new_name = f"blk.{bid}.time_mix_lerp_fused.weight"
79                data = torch.stack([self.lerp_weights[bid][f"blk.{bid}.time_mix_lerp_{i}.weight"].unsqueeze(0) for i in ["w", "k", "v", "r", "g"]], dim=0).unsqueeze(1)
80                yield (new_name, data)
81            return
82
83        yield (new_name, data_torch)
84
85
86@ModelBase.register("RWKV6Qwen2ForCausalLM")
87@ModelBase.example("recursal/QRWKV6-32B-Instruct-Preview-v0.1")
88class RWKV6Qwen2Model(Rwkv6Model):
89    model_arch = gguf.MODEL_ARCH.RWKV6QWEN2
90
91    def set_vocab(self):
92        try:
93            self._set_vocab_sentencepiece()
94        except FileNotFoundError:
95            self._set_vocab_gpt2()
96
97    def set_gguf_parameters(self):
98        num_attention_heads = self.hparams["num_attention_heads"]
99        num_key_value_heads = self.hparams["num_key_value_heads"]
100        hidden_size = self.hparams["hidden_size"]
101        head_size = hidden_size // num_attention_heads
102        rms_norm_eps = self.hparams["rms_norm_eps"]
103        intermediate_size = self.hparams["intermediate_size"]
104        time_mix_extra_dim = self.hparams.get("lora_rank_tokenshift", 64 if hidden_size >= 4096 else 32)
105        time_decay_extra_dim = self.hparams.get("lora_rank_decay", 128 if hidden_size >= 4096 else 64)
106
107        # RWKV isn't context limited
108        self.gguf_writer.add_context_length(1048576)
109        self.gguf_writer.add_embedding_length(hidden_size)
110        self.gguf_writer.add_block_count(self.block_count)
111        self.gguf_writer.add_wkv_head_size(head_size)
112        self.gguf_writer.add_time_mix_extra_dim(time_mix_extra_dim)
113        self.gguf_writer.add_time_decay_extra_dim(time_decay_extra_dim)
114        self.gguf_writer.add_feed_forward_length(intermediate_size)
115        self.gguf_writer.add_file_type(self.ftype)
116
117        # special parameters for time_mixing in RWKV6QWEN2
118        self.gguf_writer.add_layer_norm_rms_eps(rms_norm_eps)
119        self.gguf_writer.add_token_shift_count(1)
120        # RWKV6QWEN2 use grouped key/value like GQA
121        self.gguf_writer.add_head_count_kv(num_key_value_heads)
122
123        # required by llama.cpp, unused
124        self.gguf_writer.add_head_count(0)
125
126    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
127        for new_name, data in super().modify_tensors(data_torch, name, bid):
128            if "time_mix_w1" in new_name or "time_mix_w2" in new_name:
129                data = data.view(5, -1, data.shape[-1])
130                # rwkv6qwen2 has a different order of rkvwg instead of the original wkvrg
131                # permute them here to avoid code changes
132                data = torch.stack([data[3], data[1], data[2], data[0], data[4]], dim=0).view(-1, data.shape[-1])
133                if "w2" in new_name:
134                    data = data.view(5, -1, data.shape[-1])
135                yield (new_name, data)
136                continue
137            yield (new_name, data)
138
139
140@ModelBase.register("Rwkv7ForCausalLM", "RWKV7ForCausalLM")
141@ModelBase.example("fla-hub/rwkv7-1.5B-world")
142class Rwkv7Model(TextModel):
143    model_arch = gguf.MODEL_ARCH.RWKV7
144
145    def set_vocab(self):
146        self._set_vocab_rwkv_world()
147
148    def calc_lora_rank(self, hidden_size, exponent, multiplier):
149        return max(1, round(hidden_size ** exponent * multiplier / 32)) * 32
150
151    def set_gguf_parameters(self):
152        try:
153            head_size = self.hparams["head_size"]
154            layer_norm_eps = self.hparams["layer_norm_epsilon"]
155        except KeyError:
156            head_size = self.hparams["head_dim"]
157            layer_norm_eps = self.hparams["norm_eps"]
158        hidden_size = self.hparams["hidden_size"]
159        intermediate_size = self.hparams["intermediate_size"] if self.hparams["intermediate_size"] is not None else (hidden_size * 4)
160
161        # ICLR: In-Context-Learning-Rate
162        try:
163            lora_rank_decay = self.hparams["lora_rank_decay"] if self.hparams["lora_rank_decay"] is not None else self.calc_lora_rank(hidden_size, 0.5, 1.8)
164            lora_rank_iclr = self.hparams["lora_rank_iclr"] if self.hparams["lora_rank_iclr"] is not None else self.calc_lora_rank(hidden_size, 0.5, 1.8)
165            lora_rank_value_residual_mix = self.hparams["lora_rank_value_residual_mix"] if self.hparams["lora_rank_value_residual_mix"] is not None else self.calc_lora_rank(hidden_size, 0.5, 1.3)
166            lora_rank_gate = self.hparams["lora_rank_gate"] if self.hparams["lora_rank_gate"] is not None else self.calc_lora_rank(hidden_size, 0.8, 0.6)
167        except KeyError:
168            lora_rank_decay = self.hparams["decay_low_rank_dim"] if self.hparams["decay_low_rank_dim"] is not None else self.calc_lora_rank(hidden_size, 0.5, 1.8)
169            lora_rank_iclr = self.hparams["a_low_rank_dim"] if self.hparams["a_low_rank_dim"] is not None else self.calc_lora_rank(hidden_size, 0.5, 1.8)
170            lora_rank_value_residual_mix = self.hparams["v_low_rank_dim"] if self.hparams["v_low_rank_dim"] is not None else self.calc_lora_rank(hidden_size, 0.5, 1.3)
171            lora_rank_gate = self.hparams["gate_low_rank_dim"] if self.hparams["gate_low_rank_dim"] is not None else self.calc_lora_rank(hidden_size, 0.8, 0.6)
172
173        # RWKV isn't context limited
174        self.gguf_writer.add_context_length(1048576)
175        self.gguf_writer.add_embedding_length(hidden_size)
176        self.gguf_writer.add_block_count(self.block_count)
177        self.gguf_writer.add_layer_norm_eps(layer_norm_eps)
178        self.gguf_writer.add_wkv_head_size(head_size)
179        self.gguf_writer.add_decay_lora_rank(lora_rank_decay)
180        self.gguf_writer.add_iclr_lora_rank(lora_rank_iclr)
181        self.gguf_writer.add_value_residual_mix_lora_rank(lora_rank_value_residual_mix)
182        self.gguf_writer.add_gate_lora_rank(lora_rank_gate)
183        self.gguf_writer.add_feed_forward_length(intermediate_size)
184        self.gguf_writer.add_file_type(self.ftype)
185
186        # required by llama.cpp, unused
187        self.gguf_writer.add_head_count(0)
188
189    lerp_weights: dict[int, dict[str, Tensor]] = {}
190    lora_needs_transpose: bool = True
191
192    @classmethod
193    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
194        name, gen = item
195
196        # unify tensor names here to make life easier
197        name = name.replace("blocks", "layers").replace("ffn", "feed_forward")
198        name = name.replace("self_attn", "attention").replace("attn", "attention")
199        name = name.replace("time_mixer.", "")
200
201        name = name.replace("feed_forward_norm", "ln2")
202        name = name.replace("g_norm", "ln_x")
203
204        return super().filter_tensors((name, gen))
205
206    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
207        # lora layer names in fla-hub's impl
208        if "_lora.lora" in name:
209            self.lora_needs_transpose = False
210        name = name.replace("_lora.lora.0.weight", "1.weight")
211        name = name.replace("_lora.lora.2.weight", "2.weight")
212        name = name.replace("_lora.lora.2.bias", "0.weight")
213
214        if "attention.v" in name and "value" not in self.map_tensor_name(name) and bid == 0:
215            # some models have dummy v0/v1/v2 on first layer while others don't
216            # ignore them all since they are not used
217            return
218
219        wkv_has_gate = self.hparams.get("wkv_has_gate", True)
220        lerp_list = ["r", "w", "k", "v", "a", "g"] if wkv_has_gate else ["r", "w", "k", "v", "a"]
221
222        if bid is not None and "attention.x_" in name:
223            if "attention.x_x" in name:
224                # already concatenated
225                new_name = f"blk.{bid}.time_mix_lerp_fused.weight"
226                data = data_torch.reshape(len(lerp_list), 1, 1, -1)
227                yield (new_name, data)
228            else:
229                try:
230                    self.lerp_weights[bid][name] = data_torch
231                except KeyError:
232                    self.lerp_weights[bid] = {name: data_torch}
233                if all(f"model.layers.{bid}.attention.x_{i}" in self.lerp_weights[bid].keys() for i in lerp_list):
234                    new_name = f"blk.{bid}.time_mix_lerp_fused.weight"
235                    data = torch.stack([self.lerp_weights[bid][f"model.layers.{bid}.attention.x_{i}"] for i in lerp_list], dim=0)
236                    yield (new_name, data)
237            return
238        else:
239            data_torch = data_torch.squeeze()
240            new_name = self.map_tensor_name(name)
241
242            if not (new_name.endswith(".weight") or new_name.endswith(".bias")):
243                new_name += ".weight"
244
245            if self.lora_needs_transpose and any(
246                new_name.endswith(t) for t in [
247                    "time_mix_w1.weight", "time_mix_w2.weight",
248                    "time_mix_a1.weight", "time_mix_a2.weight",
249                    "time_mix_v1.weight", "time_mix_v2.weight",
250                    "time_mix_g1.weight", "time_mix_g2.weight",
251                ]
252            ):
253                data_torch = data_torch.transpose(0, 1)
254
255            if 'r_k' in new_name:
256                data_torch = data_torch.flatten()
257
258            if bid == 0 and "time_mix_a" in new_name:
259                # dummy v0/v1/v2 on first layer
260                # easiest way to make llama happy
261                yield (new_name.replace("time_mix_a", "time_mix_v"), data_torch)
262
263            yield (new_name, data_torch)
264
265
266@ModelBase.register("RwkvHybridForCausalLM")
267@ModelBase.example("RWKV-Red-Team/ARWKV-7B-Preview-0.1")
268class ARwkv7Model(Rwkv7Model):
269    model_arch = gguf.MODEL_ARCH.ARWKV7
270
271    def set_vocab(self):
272        try:
273            self._set_vocab_sentencepiece()
274        except FileNotFoundError:
275            self._set_vocab_gpt2()
276
277    def set_gguf_parameters(self):
278        hidden_size = self.hparams["hidden_size"]
279        head_size = self.hparams["head_size"]
280        rms_norm_eps = self.hparams["rms_norm_eps"]
281        intermediate_size = self.hparams["intermediate_size"]
282        wkv_has_gate = self.hparams["wkv_has_gate"]
283        assert self.hparams["wkv_version"] == 7
284
285        # ICLR: In-Context-Learning-Rate
286        lora_rank_decay = 64
287        lora_rank_iclr = 64
288        lora_rank_value_residual_mix = 32
289        lora_rank_gate = 128 if wkv_has_gate else 0
290
291        # RWKV isn't context limited
292        self.gguf_writer.add_context_length(1048576)
293        self.gguf_writer.add_embedding_length(hidden_size)
294        self.gguf_writer.add_block_count(self.block_count)
295        self.gguf_writer.add_layer_norm_rms_eps(rms_norm_eps)
296        self.gguf_writer.add_wkv_head_size(head_size)
297        self.gguf_writer.add_decay_lora_rank(lora_rank_decay)
298        self.gguf_writer.add_iclr_lora_rank(lora_rank_iclr)
299        self.gguf_writer.add_value_residual_mix_lora_rank(lora_rank_value_residual_mix)
300        self.gguf_writer.add_gate_lora_rank(lora_rank_gate)
301        self.gguf_writer.add_feed_forward_length(intermediate_size)
302        self.gguf_writer.add_file_type(self.ftype)
303        self.gguf_writer.add_token_shift_count(1)
304
305        # required by llama.cpp, unused
306        self.gguf_writer.add_head_count(0)