Brunobkr/llama.cpp_AlgMor24_github

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