Felipe97/llama-cpp-compiled

1from __future__ import annotations2 3import re4 5from typing import Callable, Iterable, TYPE_CHECKING6 7import torch8 9if TYPE_CHECKING:10    from torch import Tensor11 12from .base import ModelBase, TextModel, gguf, logger13 14from .deepseek import DeepseekV2Model15 16 17@ModelBase.register("Glm4ForCausalLM", "Glm4vForConditionalGeneration")18@ModelBase.example("zai-org/GLM-4-9B-0414")19class Glm4Model(TextModel):20    model_arch = gguf.MODEL_ARCH.GLM421    use_mrope = False22    partial_rotary_factor = 0.523 24    def __init__(self, *args, **kwargs):25        super().__init__(*args, **kwargs)26        self.partial_rotary_factor = self.rope_parameters.get("partial_rotary_factor", 0.5)27        if "mrope_section" in self.rope_parameters:28            self.use_mrope = True29            logger.info("Q/K weight will need to be permuted for M-RoPE")30 31    def set_vocab(self):32        from transformers import AutoTokenizer33        tokenizer = AutoTokenizer.from_pretrained(self.dir_model, trust_remote_code=True)34        special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True)35        tokens, toktypes, tokpre = self.get_vocab_base()36        self.gguf_writer.add_tokenizer_model("gpt2")37        self.gguf_writer.add_tokenizer_pre(tokpre)38        self.gguf_writer.add_token_list(tokens)39        self.gguf_writer.add_token_types(toktypes)40        special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True)41        special_vocab._set_special_token("eos", tokenizer.get_added_vocab()["<|endoftext|>"])  # ty: ignore[unresolved-attribute]42        special_vocab._set_special_token("eot", tokenizer.get_added_vocab()["<|user|>"])  # ty: ignore[unresolved-attribute]43        special_vocab._set_special_token("unk", tokenizer.get_added_vocab()["<|endoftext|>"])  # ty: ignore[unresolved-attribute]44        special_vocab._set_special_token("bos", tokenizer.get_added_vocab()["<|endoftext|>"])  # ty: ignore[unresolved-attribute]45        special_vocab.add_to_gguf(self.gguf_writer)46 47    def set_gguf_parameters(self):48        super().set_gguf_parameters()49        if (rope_dim := self.hparams.get("head_dim")) is None:50            rope_dim = self.hparams["hidden_size"] // self.hparams["num_attention_heads"]51        self.gguf_writer.add_rope_dimension_count(int(rope_dim * self.partial_rotary_factor))52 53    @staticmethod54    def normal_to_neox(weights: Tensor, n_head: int, n_head_kv: int, head_dim: int, partial_rotary_factor: float) -> Tensor:55        orig_shape = weights.shape56        if len(orig_shape) == 1:57            weights = weights.unsqueeze(1)  # [out_dim, 1]58        if len(weights.shape) != 2:59            raise ValueError("Only 1D and 2D tensors are supported.")60        n_effective_heads = weights.shape[0] // head_dim61        if n_head_kv is not None and n_effective_heads != n_head:62            if n_effective_heads != n_head_kv:63                raise AssertionError(f"Mismatch in effective heads: computed {n_effective_heads}, expected {n_head} or {n_head_kv}")64        rotary_dim = int(head_dim * partial_rotary_factor)65        if rotary_dim % 2 != 0:66            raise ValueError("rotary_dim must be even.")67        reshaped = weights.reshape(n_effective_heads, head_dim, -1)68        rot_part = reshaped[:, :rotary_dim, :]69        non_rot_part = reshaped[:, rotary_dim:, :]70        permuted_rot = torch.cat((rot_part[:, ::2, :], rot_part[:, 1::2, :]), dim=1)71        combined = torch.cat((permuted_rot, non_rot_part), dim=1)72        result = combined.reshape(weights.shape)73        return result if len(orig_shape) != 1 else result.squeeze(1)74 75    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:76        if self.use_mrope:77            n_head = self.hparams["num_attention_heads"]78            n_kv_head = self.hparams["num_key_value_heads"]79            n_embd = self.hparams["hidden_size"]80            head_dim = self.hparams.get("head_dim", n_embd // n_head)81            # because llama.cpp M-RoPE kernel only supports Neox ordering, we have to permute the weights here82            if name.endswith(("q_proj.weight", "q_proj.bias")):83                data_torch = Glm4Model.normal_to_neox(data_torch, n_head, n_head, head_dim, self.partial_rotary_factor)84            if name.endswith(("k_proj.weight", "k_proj.bias")):85                data_torch = Glm4Model.normal_to_neox(data_torch, n_head, n_kv_head, head_dim, self.partial_rotary_factor)86        yield from super().modify_tensors(data_torch, name, bid)87 88 89@ModelBase.register("GlmOcrForConditionalGeneration")90@ModelBase.example("zai-org/GLM-OCR")91class GlmOCRModel(Glm4Model):92    model_arch = gguf.MODEL_ARCH.GLM493    use_mrope = False94    partial_rotary_factor = 0.595 96    # Note: GLM-OCR is the same as GLM4, but with an extra NextN/MTP prediction layer97 98    def __init__(self, *args, **kwargs):99        super().__init__(*args, **kwargs)100        # GLM-OCR has num_hidden_layers + 1 actual layers (including NextN layer)101        self.block_count = self.hparams["num_hidden_layers"] + self.hparams.get("num_nextn_predict_layers", 0)102        self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)103 104    def set_gguf_parameters(self):105        super().set_gguf_parameters()106        # NextN/MTP prediction layers107        if (num_nextn_predict_layers := self.hparams.get("num_nextn_predict_layers")) is not None:108            self.gguf_writer.add_nextn_predict_layers(num_nextn_predict_layers)109 110 111@ModelBase.register("Glm4MoeForCausalLM", "Glm4vMoeForConditionalGeneration")112@ModelBase.example("zai-org/GLM-4.5-Air")113class Glm4MoeModel(TextModel):114    model_arch = gguf.MODEL_ARCH.GLM4_MOE115    supports_mtp_export = True116    _n_main_layers: int | None = None117 118    def __init__(self, *args, **kwargs):119        super().__init__(*args, **kwargs)120        if not self.no_mtp:121            self.block_count += self.hparams.get("num_nextn_predict_layers", 0)122            self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)123 124    def index_tensors(self, remote_hf_model_id: str | None = None):125        hparams = {**self.hparams, **self.hparams.get("text_config", {})}126        key = next((k for k in ["n_layers", "num_hidden_layers", "n_layer", "num_layers"] if k in hparams), None)127        type(self)._n_main_layers = hparams.get(key)128        return super().index_tensors(remote_hf_model_id=remote_hf_model_id)129 130    @classmethod131    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:132        if (titem := super().filter_tensors(item)) is None:133            return None134        name, gen = titem135 136        assert cls._n_main_layers is not None137        is_mtp = (m := re.match(r"model\.layers\.(\d+)\.", name)) is not None and int(m.group(1)) >= cls._n_main_layers138 139        if is_mtp and cls.no_mtp:140            return None141        if cls.mtp_only and not is_mtp and name not in (142            "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight",143        ):144            return None145 146        return name, gen147 148    def set_vocab(self):149        return self._set_vocab_glm()150 151    def set_gguf_parameters(self):152        super().set_gguf_parameters()153        if (rope_dim := self.hparams.get("head_dim")) is None:154            rope_dim = (155                self.hparams["hidden_size"] // self.hparams["num_attention_heads"]156            )157        self.gguf_writer.add_rope_dimension_count(158            int(rope_dim * self.rope_parameters.get("partial_rotary_factor", 0.5))159        )160 161        # MoE parameters - Use only routed expert count (shared experts handled separately)162        if (n_routed_experts := self.hparams.get("n_routed_experts")) is not None:163            self.gguf_writer.add_expert_count(n_routed_experts)164        if (moe_intermediate_size := self.hparams.get("moe_intermediate_size")) is not None:165            self.gguf_writer.add_expert_feed_forward_length(moe_intermediate_size)166        if (n_shared_experts := self.hparams.get("n_shared_experts")) is not None:167            self.gguf_writer.add_expert_shared_count(n_shared_experts)168        if (first_k_dense_replace := self.hparams.get("first_k_dense_replace")) is not None:169            self.gguf_writer.add_leading_dense_block_count(first_k_dense_replace)170 171        # Expert gating function (sigmoid for GLM4_MOE)172        self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SIGMOID)173 174        # Routed scaling factor175        if (routed_scaling_factor := self.hparams.get("routed_scaling_factor")) is not None:176            self.gguf_writer.add_expert_weights_scale(routed_scaling_factor)177 178        # Normalise topk probabilities179        if (norm_topk_prob := self.hparams.get("norm_topk_prob")) is not None:180            self.gguf_writer.add_expert_weights_norm(norm_topk_prob)181 182        if not self.no_mtp and (num_nextn_predict_layers := self.hparams.get("num_nextn_predict_layers")) is not None:183            self.gguf_writer.add_nextn_predict_layers(num_nextn_predict_layers)184 185    def prepare_metadata(self, vocab_only: bool):186        from_dir = self.fname_out.is_dir()187        super().prepare_metadata(vocab_only=vocab_only)188 189        if not self.mtp_only or not from_dir:190            return191 192        output_type: str = self.ftype.name.partition("_")[2]193        fname_default: str = gguf.naming_convention(194            self.metadata.name, self.metadata.basename, self.metadata.finetune,195            self.metadata.version, size_label=None, output_type=output_type, model_type=None)196        self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf"197 198    _experts: list[dict[str, Tensor]] | None = None199 200    # note: unlike GLM4V non-MoE, we don't need to permute Q/K here since GLM4V_MOE uses Neox ordering already201    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:202        # Handle main token embedding (but not layer-specific NextN embeddings)203        if name == "model.embed_tokens.weight" and ".layers." not in name:204            yield from super().modify_tensors(data_torch, "token_embd.weight", bid)205            return206 207        # Handle routed experts208        if name.find("mlp.experts") != -1:209            n_experts = self.hparams["n_routed_experts"]210            assert bid is not None211 212            if self._experts is None:213                self._experts = [{} for _ in range(self.block_count)]214 215            self._experts[bid][name] = data_torch216 217            if len(self._experts[bid]) >= n_experts * 3:218                # merge the experts into a single 3d tensor219                for w_name in ["down_proj", "gate_proj", "up_proj"]:220                    datas: list[Tensor] = []221 222                    for xid in range(n_experts):223                        ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"224                        datas.append(self._experts[bid][ename])225                        del self._experts[bid][ename]226 227                    data_torch = torch.stack(datas, dim=0)228 229                    merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"230 231                    yield from super().modify_tensors(data_torch, merged_name, bid)232                return233            else:234                return235 236        yield from super().modify_tensors(data_torch, name, bid)237 238    def prepare_tensors(self):239        super().prepare_tensors()240        if self._experts is not None:241            # flatten `list[dict[str, Tensor]]` into `list[str]`242            experts = [k for d in self._experts for k in d.keys()]243            if len(experts) > 0:244                raise ValueError(f"Unprocessed experts: {experts}")245 246 247@ModelBase.register("Glm4MoeLiteForCausalLM")248@ModelBase.example("zai-org/GLM-4.7-Flash")249class Glm4MoeLiteModel(DeepseekV2Model):250    model_arch = gguf.MODEL_ARCH.DEEPSEEK2251    skip_mtp = False252    supports_mtp_export = True253    _n_main_layers: int | None = None254 255    def set_vocab(self):256        return self._set_vocab_glm()257 258    def __init__(self, *args, **kwargs):259        super().__init__(*args, **kwargs)260 261        num_hidden_layers = self.hparams["num_hidden_layers"]262        self.num_nextn_predict_layers = self.hparams.get("num_nextn_predict_layers", 0)263        self.skip_mtp = self.no_mtp or self.num_nextn_predict_layers == 0264 265        if self.skip_mtp:266            self.block_count = num_hidden_layers267        else:268            self.block_count = num_hidden_layers + self.num_nextn_predict_layers269 270        self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)271 272    def set_gguf_parameters(self):273        super().set_gguf_parameters()274 275        if self.skip_mtp:276            return277 278        self.gguf_writer.add_nextn_predict_layers(self.num_nextn_predict_layers)279 280    def index_tensors(self, remote_hf_model_id: str | None = None):281        type(self)._n_main_layers = self.hparams["num_hidden_layers"]282        return super().index_tensors(remote_hf_model_id=remote_hf_model_id)283 284    @classmethod285    def filter_tensors(cls, item):286        if (titem := super().filter_tensors(item)) is None:287            return None288        name, gen = titem289 290        if cls._n_main_layers is not None:291            match = re.match(r"model\.layers\.(\d+)\.", name)292            is_mtp = match is not None and int(match.group(1)) >= cls._n_main_layers293            if is_mtp and cls.no_mtp:294                return None295            if cls.mtp_only and not is_mtp and name not in (296                "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight",297            ):298                return None299 300        return name, gen301 302    def prepare_metadata(self, vocab_only: bool):303        from_dir = self.fname_out.is_dir()304        super().prepare_metadata(vocab_only=vocab_only)305 306        if not self.mtp_only or not from_dir:307            return308 309        output_type: str = self.ftype.name.partition("_")[2]310        fname_default: str = gguf.naming_convention(311            self.metadata.name, self.metadata.basename, self.metadata.finetune,312            self.metadata.version, size_label=None, output_type=output_type, model_type=None)313        self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf"314 315 316@ModelBase.register("GlmMoeDsaForCausalLM")317@ModelBase.example("zai-org/GLM-5.2")318class GlmMoeDsaModel(DeepseekV2Model):319    model_arch = gguf.MODEL_ARCH.GLM_DSA320    skip_mtp = False321    supports_mtp_export = True322 323    # Trunk layer count, stashed before indexing so the classmethod324    # filter_tensors can identify the appended NextN/MTP block (mirrors325    # HYV3Model / Step35Model).326    _n_main_layers: int | None = None327 328    def __init__(self, *args, **kwargs):329        super().__init__(*args, **kwargs)330        self.block_count = self.hparams["num_hidden_layers"]331        if not self.no_mtp:332            self.block_count += self.hparams.get("num_nextn_predict_layers", 0)333        self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)334 335    def index_tensors(self, remote_hf_model_id: str | None = None):336        type(self)._n_main_layers = self.hparams["num_hidden_layers"]337        return super().index_tensors(remote_hf_model_id=remote_hf_model_id)338 339    @classmethod340    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:341        if (titem := super().filter_tensors(item)) is None:342            return None343        name, gen = titem344 345        # GLM-5.2 appends the NextN/MTP block past num_hidden_layers346        # (model.layers.78 -> blk.78 in the 79-block file).347        assert cls._n_main_layers is not None348        is_mtp = (m := re.match(r"model\.layers\.(\d+)\.", name)) is not None and int(m.group(1)) >= cls._n_main_layers349 350        # --no-mtp: drop the appended NextN block entirely.351        if is_mtp and cls.no_mtp:352            return None353        # --mtp: keep ONLY NextN-block tensors plus the shared embeddings/354        # norm/lm_head (so the resulting GGUF carries just the draft head).355        if cls.mtp_only and not is_mtp and name not in (356            "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight",357        ):358            return None359 360        return name, gen361 362    def set_vocab(self):363        return self._set_vocab_glm()364 365    def set_gguf_parameters(self):366        super().set_gguf_parameters()367 368        rope_dim = self.hparams["qk_rope_head_dim"]369        partial_rotary_factor = self.rope_parameters.get("partial_rotary_factor", 1.0)370        self.gguf_writer.add_rope_dimension_count(int(rope_dim * partial_rotary_factor))371 372        # NextN/MTP prediction layers373        if not self.no_mtp and (num_nextn_predict_layers := self.hparams.get("num_nextn_predict_layers")) is not None:374            self.gguf_writer.add_nextn_predict_layers(num_nextn_predict_layers)375 376        # DSA indexer parameters377        self.gguf_writer.add_indexer_head_count(self.hparams["index_n_heads"])378        self.gguf_writer.add_indexer_key_length(self.hparams["index_head_dim"])379        self.gguf_writer.add_indexer_top_k(self.hparams["index_topk"])380        if (indexer_types := self.hparams.get("indexer_types")) is not None:381            indexer_types = [t == "full" for t in indexer_types]382            self.gguf_writer.add_indexer_types(indexer_types)383 384 385@ModelBase.register("SolarOpenForCausalLM")386@ModelBase.example("upstage/Solar-Open-100B")387class SolarOpenModel(Glm4MoeModel):388    model_arch = gguf.MODEL_ARCH.GLM4_MOE389    supports_mtp_export = False390 391    def set_vocab(self):392        from transformers import AutoTokenizer393        tokenizer = AutoTokenizer.from_pretrained(self.dir_model)394        special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True)395        tokens, toktypes, tokpre = self.get_vocab_base()396        self.gguf_writer.add_tokenizer_model("gpt2")397        self.gguf_writer.add_tokenizer_pre(tokpre)398        self.gguf_writer.add_token_list(tokens)399        self.gguf_writer.add_token_types(toktypes)400        special_vocab._set_special_token("eos", tokenizer.get_added_vocab()["<|endoftext|>"])  # ty: ignore[unresolved-attribute]401        special_vocab._set_special_token("eot", tokenizer.get_added_vocab()["<|endoftext|>"])  # ty: ignore[unresolved-attribute]402        special_vocab._set_special_token("unk", tokenizer.get_added_vocab()["<unk>"])  # ty: ignore[unresolved-attribute]403        special_vocab._set_special_token("bos", tokenizer.get_added_vocab()["<|startoftext|>"])  # ty: ignore[unresolved-attribute]404        special_vocab.add_to_gguf(self.gguf_writer)405