Felipe97/llama-cpp-compiled

1from __future__ import annotations2 3from typing import Callable, Iterable, TYPE_CHECKING4 5import torch6 7if TYPE_CHECKING:8    from torch import Tensor9 10from .base import ModelBase, TextModel, gguf11 12 13@ModelBase.register("BailingMoeForCausalLM")14@ModelBase.example("inclusionAI/Ling-lite")15class BailingMoeModel(TextModel):16    model_arch = gguf.MODEL_ARCH.BAILINGMOE17 18    def set_vocab(self):19        self._set_vocab_gpt2()20 21    def set_gguf_parameters(self):22        super().set_gguf_parameters()23        hparams = self.hparams24        if (rope_dim := hparams.get("head_dim")) is None:25            rope_dim = hparams["hidden_size"] // hparams["num_attention_heads"]26 27        self.gguf_writer.add_rope_dimension_count(rope_dim)28        self.gguf_writer.add_leading_dense_block_count(hparams["first_k_dense_replace"])29        self.gguf_writer.add_vocab_size(hparams["vocab_size"])30        self.gguf_writer.add_expert_feed_forward_length(hparams["moe_intermediate_size"])31        self.gguf_writer.add_expert_weights_scale(1.0)32        self.gguf_writer.add_expert_shared_count(hparams["num_shared_experts"])33        self.gguf_writer.add_expert_weights_norm(hparams["norm_topk_prob"])34 35    _experts: list[dict[str, Tensor]] | None = None36 37    @staticmethod38    def permute(weights: Tensor, n_head: int, n_head_kv: int | None):39        if n_head_kv is not None and n_head != n_head_kv:40            n_head = n_head_kv41        return (weights.reshape(n_head, 2, weights.shape[0] // n_head // 2, *weights.shape[1:])42                .swapaxes(1, 2)43                .reshape(weights.shape))44 45    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:46        n_head = self.hparams["num_attention_heads"]47        n_kv_head = self.hparams.get("num_key_value_heads")48        n_embd = self.hparams["hidden_size"]49        if (head_dim := self.hparams.get("head_dim")) is None:50            head_dim = n_embd // n_head51 52        output_name = self.format_tensor_name(gguf.MODEL_TENSOR.OUTPUT)53 54        if name.endswith("attention.dense.weight"):55            yield from super().modify_tensors(data_torch, self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_OUT, bid), bid)56            return57        elif name.endswith("query_key_value.weight"):58            q, k, v = data_torch.split([n_head * head_dim, n_kv_head * head_dim, n_kv_head * head_dim], dim=-2)59 60            yield from super().modify_tensors(BailingMoeModel.permute(q, n_head, n_head), self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_Q, bid), bid)61            yield from super().modify_tensors(BailingMoeModel.permute(k, n_head, n_kv_head), self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_K, bid), bid)62            yield from super().modify_tensors(v,self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_V, bid), bid)63            return64        elif name.find("mlp.experts") != -1:65            n_experts = self.find_hparam(["num_local_experts", "num_experts"])66            assert bid is not None67 68            if self._experts is None:69                self._experts = [{} for _ in range(self.block_count)]70 71            self._experts[bid][name] = data_torch72 73            if len(self._experts[bid]) >= n_experts * 3:74                # merge the experts into a single 3d tensor75                for w_name in ["down_proj", "gate_proj", "up_proj"]:76                    datas: list[Tensor] = []77 78                    for xid in range(n_experts):79                        ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"80                        datas.append(self._experts[bid][ename])81                        del self._experts[bid][ename]82 83                    data_torch = torch.stack(datas, dim=0)84 85                    merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"86 87                    new_name = self.map_tensor_name(merged_name)88 89                    yield from super().modify_tensors(data_torch, new_name, bid)90 91            return92 93        new_name = self.map_tensor_name(name)94 95        if new_name == output_name and self.hparams.get("norm_head"):96            data_torch = data_torch.float()97            data_torch /= torch.norm(data_torch, p=2, dim=0, keepdim=True) + 1e-798 99        yield from super().modify_tensors(data_torch, new_name, bid)100 101    def prepare_tensors(self):102        super().prepare_tensors()103 104        if self._experts is not None:105            # flatten `list[dict[str, Tensor]]` into `list[str]`106            experts = [k for d in self._experts for k in d.keys()]107            if len(experts) > 0:108                raise ValueError(f"Unprocessed experts: {experts}")109 110 111@ModelBase.register("BailingMoeV2ForCausalLM")112@ModelBase.example("inclusionAI/Ling-mini-2.0")113class BailingMoeV2Model(TextModel):114    model_arch = gguf.MODEL_ARCH.BAILINGMOE2115 116    def __init__(self, *args, **kwargs):117        super().__init__(*args, **kwargs)118        if nextn_layers := self.hparams.get("num_nextn_predict_layers", 0):119            self.block_count = self.hparams["num_hidden_layers"] + nextn_layers120            self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)121 122    def set_vocab(self):123        self._set_vocab_gpt2()124 125    def set_gguf_parameters(self):126        super().set_gguf_parameters()127        hparams = self.hparams128        if (rope_dim := hparams.get("head_dim")) is None:129            rope_dim = hparams["hidden_size"] // hparams["num_attention_heads"]130 131        self.gguf_writer.add_rope_dimension_count(int(rope_dim * self.rope_parameters.get("partial_rotary_factor", 0.5)))132        self.gguf_writer.add_leading_dense_block_count(hparams["first_k_dense_replace"])133        self.gguf_writer.add_vocab_size(hparams["vocab_size"])134        self.gguf_writer.add_expert_feed_forward_length(hparams["moe_intermediate_size"])135        self.gguf_writer.add_expert_shared_feed_forward_length(hparams.get("moe_shared_expert_intermediate_size", hparams["moe_intermediate_size"] * hparams["num_shared_experts"]))136        self.gguf_writer.add_expert_weights_scale(hparams["routed_scaling_factor"])137        self.gguf_writer.add_expert_shared_count(hparams["num_shared_experts"])138        self.gguf_writer.add_expert_weights_norm(hparams["norm_topk_prob"])139 140        if (nextn_layers := self.hparams.get("num_nextn_predict_layers")) is not None:141            self.gguf_writer.add_nextn_predict_layers(nextn_layers)142 143    _experts: list[dict[str, Tensor]] | None = None144 145    @classmethod146    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:147        name, gen = item148 149        if name.endswith(".expert_bias"):150            name = name.replace(".expert_bias", ".expert_bias.bias")151 152        return super().filter_tensors((name, gen))153 154    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:155        if "mlp.experts" in name:156            n_experts = self.find_hparam(["num_local_experts", "num_experts"])157            assert bid is not None158 159            if self._experts is None:160                self._experts = [{} for _ in range(self.block_count)]161 162            self._experts[bid][name] = data_torch163 164            if len(self._experts[bid]) >= n_experts * 3:165                # merge the experts into a single 3d tensor166                for w_name in ["down_proj", "gate_proj", "up_proj"]:167                    datas: list[Tensor] = []168 169                    for xid in range(n_experts):170                        ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"171                        datas.append(self._experts[bid][ename])172                        del self._experts[bid][ename]173 174                    data_torch = torch.stack(datas, dim=0)175 176                    merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"177 178                    yield from super().modify_tensors(data_torch, merged_name, bid)179            return180 181        yield from super().modify_tensors(data_torch, name, bid)182 183    def prepare_tensors(self):184        super().prepare_tensors()185 186        if self._experts is not None:187            # flatten `list[dict[str, Tensor]]` into `list[str]`188            experts = [k for d in self._experts for k in d.keys()]189            if len(experts) > 0:190                raise ValueError(f"Unprocessed experts: {experts}")191 192 193@ModelBase.register("SarvamMoEForCausalLM", "modeling_sarvam_moe.SarvamMoEForCausalLM")194@ModelBase.example("sarvamai/sarvam-30b")195class SarvamMoEModel(BailingMoeV2Model):196    model_arch = gguf.MODEL_ARCH.BAILINGMOE2197    # Sarvam-MoE shares the BailingMoeV2 architecture; only differences:198    #  - full rotary (no partial_rotary_factor)199    #  - expert bias is zero-mean normalized at load time200 201    def set_gguf_parameters(self):202        super().set_gguf_parameters()203        hparams = self.hparams204        if (rope_dim := hparams.get("head_dim")) is None:205            rope_dim = hparams["hidden_size"] // hparams["num_attention_heads"]206        # Override the partial-rotary value written by BailingMoeV2 with the full rotary dim207        self.gguf_writer.add_rope_dimension_count(rope_dim)208 209    @classmethod210    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:211        name, gen = item212        if name.endswith(".expert_bias"):213            # Sarvam normalizes expert bias to zero mean214            inner = gen215 216            def gen():217                t = inner()218                return t - t.mean()219        return super().filter_tensors((name, gen))220