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
14.9 KiB312 linesraw
1from __future__ import annotations
2
3import math
4
5from pathlib import Path
6from typing import Callable, Iterable, TYPE_CHECKING
7
8import torch
9
10if TYPE_CHECKING:
11    from torch import Tensor
12
13from .base import MmprojModel, ModelBase, TextModel, gguf
14from .qwenvl import Qwen2VLVisionModel
15
16
17@ModelBase.register("ExaoneForCausalLM")
18@ModelBase.example("LGAI-EXAONE/EXAONE-3.5-2.4B-Instruct")
19class ExaoneModel(TextModel):
20    model_arch = gguf.MODEL_ARCH.EXAONE
21
22    def set_gguf_parameters(self):
23        super().set_gguf_parameters()
24        hparams = self.hparams
25
26        assert (hparams["activation_function"] == "silu")
27
28        rotary_factor = self.rope_parameters.get("partial_rotary_factor")
29        rotary_factor = rotary_factor if rotary_factor is not None else 1.0
30        self.gguf_writer.add_rope_dimension_count(int(rotary_factor * (hparams["hidden_size"] // hparams["num_attention_heads"])))
31
32    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
33        if rope_params := self.rope_parameters.get("full_attention", self.rope_parameters):
34            if rope_params.get("rope_type", '').lower() == "llama3":
35                base = self.rope_parameters.get("rope_theta", 10000.0)
36                if (dim := self.hparams.get("head_dim")) is None:
37                    dim = self.hparams["hidden_size"] // self.hparams["num_attention_heads"]
38                freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
39
40                factor = rope_params.get("factor", 8.0)
41                low_freq_factor = rope_params.get("low_freq_factor", 1.0)
42                high_freq_factor = rope_params.get("high_freq_factor", 4.0)
43                old_context_len = rope_params.get("original_max_position_embeddings", 8192)
44
45                low_freq_wavelen = old_context_len / low_freq_factor
46                high_freq_wavelen = old_context_len / high_freq_factor
47                assert low_freq_wavelen != high_freq_wavelen
48
49                rope_factors = []
50                for freq in freqs:
51                    wavelen = 2 * math.pi / freq
52                    if wavelen < high_freq_wavelen:
53                        rope_factors.append(1)
54                    elif wavelen > low_freq_wavelen:
55                        rope_factors.append(factor)
56                    else:
57                        smooth = (old_context_len / wavelen - low_freq_factor) / (high_freq_factor - low_freq_factor)
58                        rope_factors.append(1 / ((1 - smooth) / factor + smooth))
59
60                yield (self.format_tensor_name(gguf.MODEL_TENSOR.ROPE_FREQS), torch.tensor(rope_factors, dtype=torch.float32))
61
62
63@ModelBase.register("Exaone4ForCausalLM")
64@ModelBase.example("LGAI-EXAONE/EXAONE-4.0-32B")
65class Exaone4Model(TextModel):
66    model_arch = gguf.MODEL_ARCH.EXAONE4
67
68    def set_vocab(self):
69        tokens, toktypes, tokpre = self.get_vocab_base()
70        self.gguf_writer.add_tokenizer_model("gpt2")
71        self.gguf_writer.add_tokenizer_pre(tokpre)
72        self.gguf_writer.add_token_list(tokens)
73        self.gguf_writer.add_token_types(toktypes)
74
75        special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True)
76        special_vocab.add_to_gguf(self.gguf_writer)
77
78    def set_gguf_parameters(self):
79        super().set_gguf_parameters()
80        hparams = self.hparams
81        self.gguf_writer.add_vocab_size(hparams["vocab_size"])
82
83        if hparams.get("sliding_window") is not None:
84            self.gguf_writer.add_sliding_window(hparams["sliding_window"])
85            if "layer_types" in hparams:
86                self.gguf_writer.add_sliding_window_pattern([t == "sliding_attention" for t in hparams["layer_types"]])
87            elif "sliding_window_pattern" in hparams:
88                sliding_window_pattern = []
89                if isinstance(hparams["sliding_window_pattern"], str):  # e.g. LLLG
90                    for i in range(hparams["num_hidden_layers"]):
91                        sliding_window_pattern.append(hparams["sliding_window_pattern"][i % len(hparams["sliding_window_pattern"])] == "L")
92                if isinstance(hparams["sliding_window_pattern"], int):  # e.g. 4
93                    for i in range(hparams["num_hidden_layers"]):
94                        sliding_window_pattern.append((i + 1) % hparams["sliding_window_pattern"] != 0)
95                if len(sliding_window_pattern) == hparams["num_hidden_layers"]:
96                    self.gguf_writer.add_sliding_window_pattern(sliding_window_pattern)
97
98    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
99        if rope_params := self.rope_parameters.get("full_attention", self.rope_parameters):
100            if rope_params.get("rope_type", '').lower() == "llama3":
101                base = rope_params.get("rope_theta", 10_000.0)
102                if (dim := self.hparams.get("head_dim")) is None:
103                    dim = self.hparams["hidden_size"] // self.hparams["num_attention_heads"]
104                freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
105
106                factor = rope_params.get("factor", 16.0)
107                low_freq_factor = rope_params.get("low_freq_factor", 1.0)
108                high_freq_factor = rope_params.get("high_freq_factor", 4.0)
109                old_context_len = rope_params.get("original_max_position_embeddings", 8192)
110
111                low_freq_wavelen = old_context_len / low_freq_factor
112                high_freq_wavelen = old_context_len / high_freq_factor
113
114                rope_factors = []
115                for freq in freqs:
116                    wavelen = 2 * math.pi / freq
117                    if wavelen < high_freq_wavelen:
118                        rope_factors.append(1)
119                    elif wavelen > low_freq_wavelen:
120                        rope_factors.append(factor)
121                    else:
122                        smooth = (old_context_len / wavelen - low_freq_factor) / (high_freq_factor - low_freq_factor)
123                        rope_factors.append(1 / ((1 - smooth) / factor + smooth))
124
125                yield (self.format_tensor_name(gguf.MODEL_TENSOR.ROPE_FREQS), torch.tensor(rope_factors, dtype=torch.float32))
126
127
128# note: transformers >= 5.1 renamed the class to "ExaoneMoeForCausalLM" (lowercase 'e'),
129#       so accept both spellings - LG AI have updated the configs of already-released models
130@ModelBase.register("ExaoneMoEForCausalLM", "ExaoneMoeForCausalLM")
131@ModelBase.example("LGAI-EXAONE/K-EXAONE-236B-A23B")
132class ExaoneMoEModel(Exaone4Model):
133    model_arch = gguf.MODEL_ARCH.EXAONE_MOE
134
135    def __init__(self, *args, **kwargs):
136        super().__init__(*args, **kwargs)
137        self.block_count = self.hparams["num_hidden_layers"] + self.hparams.get("num_nextn_predict_layers", 0)
138        self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
139
140    def set_gguf_parameters(self):
141        super().set_gguf_parameters()
142        moe_intermediate_size = self.hparams["moe_intermediate_size"]
143        num_shared_experts = self.hparams["num_shared_experts"]
144        self.gguf_writer.add_expert_feed_forward_length(moe_intermediate_size)
145        self.gguf_writer.add_expert_shared_count(num_shared_experts)
146        self.gguf_writer.add_expert_shared_feed_forward_length(moe_intermediate_size * num_shared_experts)
147        self.gguf_writer.add_expert_weights_scale(self.hparams["routed_scaling_factor"])
148        self.gguf_writer.add_expert_weights_norm(self.hparams["norm_topk_prob"])
149        n_dense_layer = self.hparams.get("first_k_dense_replace", self.hparams.get("first_last_k_dense_replace", 0))
150        self.gguf_writer.add_leading_dense_block_count(n_dense_layer)
151        self.gguf_writer.add_nextn_predict_layers(self.hparams.get("num_nextn_predict_layers", 0))
152
153        self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.NONE)
154
155    _experts: list[dict[str, Tensor]] | None = None
156
157    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
158        if name.startswith("mtp."):
159            if name.find("layers.") != -1:
160                # `mtp.layers.0.[module_name]` format
161                name = name.replace(f"mtp.layers.{bid}", f"model.layers.{bid + self.hparams['num_hidden_layers']}")
162            else:
163                # mtp fc/norm weights
164                remapper = {
165                    "mtp.fc": "model.layers.{bid}.eh_proj",
166                    "mtp.pre_fc_norm_embedding": "model.layers.{bid}.enorm",
167                    "mtp.pre_fc_norm_hidden": "model.layers.{bid}.hnorm",
168                    "mtp.norm": "model.layers.{bid}.shared_head.norm",
169                }
170                _n = Path(name)
171                new_name = remapper[_n.stem] + _n.suffix
172
173                # set shared weights for all NextN/MTP layers
174                for bid in range(self.hparams['num_hidden_layers'], self.block_count):
175                    yield from super().modify_tensors(data_torch, new_name.format(bid=bid), bid)
176                return
177
178        if name.find("mlp.experts") != -1:
179            n_experts = self.find_hparam(["num_local_experts", "num_experts"])
180            assert bid is not None
181
182            if self._experts is None:
183                self._experts = [{} for _ in range(self.block_count)]
184
185            self._experts[bid][name] = data_torch
186
187            if len(self._experts[bid]) >= n_experts * 3:
188                # merge the experts into a single 3d tensor
189                for w_name in ["down_proj", "gate_proj", "up_proj"]:
190                    datas: list[Tensor] = []
191
192                    for xid in range(n_experts):
193                        ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"
194                        datas.append(self._experts[bid][ename])
195                        del self._experts[bid][ename]
196
197                    data_torch = torch.stack(datas, dim=0)
198
199                    merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"
200
201                    new_name = self.map_tensor_name(merged_name)
202
203                    yield from super().modify_tensors(data_torch, new_name, bid)
204                return
205            else:
206                return
207
208        yield from super().modify_tensors(data_torch, name, bid)
209
210    def prepare_tensors(self):
211        super().prepare_tensors()
212        if self._experts is not None:
213            # flatten `list[dict[str, Tensor]]` into `list[str]`
214            experts = [k for d in self._experts for k in d.keys()]
215            if len(experts) > 0:
216                raise ValueError(f"Unprocessed experts: {experts}")
217
218
219@ModelBase.register("Exaone4_5_ForConditionalGeneration")
220@ModelBase.example("LGAI-EXAONE/EXAONE-4.5-33B")
221class Exaone4_5_TextModel(Exaone4Model):
222    """Text tower of EXAONE 4.5; Tensors match EXAONE4"""
223
224    model_arch = gguf.MODEL_ARCH.EXAONE4
225
226    def __init__(self, *args, **kwargs):
227        super().__init__(*args, **kwargs)
228        n_nextn = int(self.hparams.get("num_nextn_predict_layers", 0) or 0)
229        if n_nextn > 0:
230            self.block_count = self.hparams["num_hidden_layers"] + n_nextn
231            self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
232
233    def set_gguf_parameters(self):
234        super().set_gguf_parameters()
235        n_nextn = int(self.hparams.get("num_nextn_predict_layers", 0) or 0)
236        if n_nextn > 0:
237            self.gguf_writer.add_nextn_predict_layers(n_nextn)
238
239    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
240        if name.startswith("mtp."):
241            n_nextn = int(self.hparams.get("num_nextn_predict_layers", 0) or 0)
242            if n_nextn <= 0:
243                return
244            nh = self.hparams["num_hidden_layers"]
245            if ".layers." in name:
246                share = self.hparams.get("mtp_share_layers", False)
247                mtp_bid = bid if bid is not None else 0
248                if share:
249                    for k in range(n_nextn):
250                        nn = name.replace(f"mtp.layers.{mtp_bid}", f"model.layers.{nh + k}")
251                        yield from super().modify_tensors(data_torch, nn, nh + k)
252                    return
253                name = name.replace(f"mtp.layers.{mtp_bid}", f"model.layers.{mtp_bid + nh}")
254            else:
255                remapper = {
256                    "mtp.fc": gguf.MODEL_TENSOR.NEXTN_EH_PROJ,
257                    "mtp.pre_fc_norm_embedding": gguf.MODEL_TENSOR.NEXTN_ENORM,
258                    "mtp.pre_fc_norm_hidden": gguf.MODEL_TENSOR.NEXTN_HNORM,
259                    "mtp.norm": gguf.MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM,
260                }
261                _n = Path(name)
262                key = _n.stem
263                if key not in remapper:
264                    return
265                for bid_mtp in range(nh, self.block_count):
266                    mapped_name = self.format_tensor_name(remapper[key], bid_mtp, suffix=_n.suffix)
267                    yield from ModelBase.modify_tensors(self, data_torch, mapped_name, bid_mtp)
268                return
269
270        yield from super().modify_tensors(data_torch, name, bid)
271
272
273@ModelBase.register("Exaone4_5_ForConditionalGeneration")
274@ModelBase.example("LGAI-EXAONE/EXAONE-4.5-33B")
275class Exaone4_5VisionModel(Qwen2VLVisionModel):
276    """Vision tower for EXAONE 4.5; Qwen2-VL-style ViT (GQA) + patch merger"""
277
278    @classmethod
279    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
280        name, gen = item
281        name = name.replace("model.visual.", "visual.", 1)
282        return super().filter_tensors((name, gen))
283
284    def set_gguf_parameters(self):
285        MmprojModel.set_gguf_parameters(self)
286        assert self.hparams_vision is not None
287        hparams = self.hparams_vision
288        self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.EXAONE4_5)
289        self.gguf_writer.add_vision_use_silu(True)
290        self.gguf_writer.add_vision_min_pixels(self.preprocessor_config["min_pixels"])
291        self.gguf_writer.add_vision_max_pixels(self.preprocessor_config["max_pixels"])
292        num_kv_head = self.find_vparam(["num_key_value_heads"], optional=True)
293        if num_kv_head is not None:
294            self.gguf_writer.add_vision_head_count_kv(num_kv_head)
295        eps = hparams.get("rms_norm_eps", self.global_config.get("rms_norm_eps", 1e-6))
296        self.gguf_writer.add_vision_attention_layernorm_eps(eps)
297        if (window_size := hparams.get("window_size")) is not None:
298            self.gguf_writer.add_vision_window_size(window_size)
299        fullatt_block_indexes = hparams.get("fullatt_block_indexes")
300        if fullatt_block_indexes:
301            n_wa_pattern = fullatt_block_indexes[0] + 1
302            for i in range(1, len(fullatt_block_indexes)):
303                if fullatt_block_indexes[i] - fullatt_block_indexes[i - 1] != n_wa_pattern:
304                    raise ValueError(f"Invalid EXAONE4.5 fullatt_block_indexes: {fullatt_block_indexes}")
305            self.gguf_writer.add_vision_n_wa_pattern(n_wa_pattern)
306
307    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
308        if ".qkv." in name:
309            yield from ModelBase.modify_tensors(self, data_torch, name, bid)
310            return
311
312        yield from Qwen2VLVisionModel.modify_tensors(self, data_torch, name, bid)