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, gguf13 14 15@ModelBase.register("BailingMoeV3ForCausalLM")16@ModelBase.example("inclusionAI/Ling-3.0-tiny", "inclusionAI/Ling-3.0-flash")17class BailingMoeV3Model(TextModel):18    model_arch = gguf.MODEL_ARCH.BAILINGMOE319    supports_mtp_export = True20 21    _experts: list[dict[str, Tensor]] | None = None22    _main_layers: int | None = None23 24    def __init__(self, *args, **kwargs):25        super().__init__(*args, **kwargs)26        nextn_layers = self.hparams.get("num_nextn_predict_layers", 0) or 027        if self.no_mtp:28            nextn_layers = 029        self.block_count = self.hparams["num_hidden_layers"] + nextn_layers30        self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)31 32    def index_tensors(self, remote_hf_model_id: str | None = None):33        type(self)._main_layers = self.hparams["num_hidden_layers"]34        return super().index_tensors(remote_hf_model_id=remote_hf_model_id)35 36    def set_vocab(self):37        self._set_vocab_gpt2()38 39    def is_full_attention(self, bid: int) -> bool:40        n_layer = self.hparams["num_hidden_layers"]41        layer_group_size = self.hparams["layer_group_size"]42        return bid >= n_layer or (bid + 1) % layer_group_size == 0 or bid >= n_layer // layer_group_size * layer_group_size43 44    def set_gguf_parameters(self):45        if not self.hparams.get("no_kda_lora", False):46            raise ValueError("BailingMoeV3 KDA LoRA projections are not supported")47        if not self.hparams.get("kda_safe_gate", False):48            raise ValueError("BailingMoeV3 non-safe KDA gates are not supported")49        if self.hparams.get("gated_attention_proj_granularity_type") != "head_wise":50            raise ValueError("BailingMoeV3 requires head-wise attention gates")51 52        self.hparams["num_key_value_heads"] = 153        super().set_gguf_parameters()54 55        n_head_kv = [1 if self.is_full_attention(il) else 0 for il in range(self.block_count)]56        self.gguf_writer.add_head_count_kv(n_head_kv)57 58        self.gguf_writer.add_vocab_size(self.hparams["vocab_size"])59        self.gguf_writer.add_ssm_conv_kernel(self.hparams["short_conv_kernel_size"])60        self.gguf_writer.add_kda_head_dim(self.hparams["head_dim"])61        self.gguf_writer.add_kda_safe_gate(self.hparams["kda_safe_gate"])62        self.gguf_writer.add_kda_gate_lower_bound(self.hparams["kda_lower_bound"])63 64        kv_lora_rank = self.hparams["kv_lora_rank"]65        qk_nope_head_dim = self.hparams["qk_nope_head_dim"]66        qk_rope_head_dim = self.hparams["qk_rope_head_dim"]67        if (q_lora_rank := self.hparams.get("q_lora_rank")) is not None:68            self.gguf_writer.add_q_lora_rank(q_lora_rank)69        self.gguf_writer.add_kv_lora_rank(kv_lora_rank)70        self.gguf_writer.add_rope_dimension_count(qk_rope_head_dim)71        self.gguf_writer.add_key_length(kv_lora_rank + qk_rope_head_dim)72        self.gguf_writer.add_key_length_mla(qk_nope_head_dim + qk_rope_head_dim)73        self.gguf_writer.add_value_length_mla(self.hparams["v_head_dim"])74 75        self.gguf_writer.add_expert_feed_forward_length(self.hparams["moe_intermediate_size"])76        self.gguf_writer.add_expert_shared_feed_forward_length(self.hparams["moe_shared_expert_intermediate_size"])77        self.gguf_writer.add_expert_shared_count(self.hparams["num_shared_experts"])78        self.gguf_writer.add_leading_dense_block_count(self.hparams["first_k_dense_replace"])79        self.gguf_writer.add_expert_weights_scale(self.hparams["routed_scaling_factor"])80        self.gguf_writer.add_expert_weights_norm(self.hparams["norm_topk_prob"])81 82        def clamp_limits(key: str) -> list[float] | None:83            values = self.hparams.get(key)84            if values is None:85                return None86            values = [0.0 if value is None else float(value) for value in values[:self.block_count]]87            return values + [0.0] * (self.block_count - len(values))88 89        if (values := clamp_limits("expert_swiglu_limit_list")) is not None:90            self.gguf_writer.add_swiglu_clamp_exp(values)91        if (values := clamp_limits("share_expert_swiglu_limit_list")) is not None:92            self.gguf_writer.add_swiglu_clamp_shexp(values)93 94        if not self.no_mtp and (nextn_layers := self.hparams.get("num_nextn_predict_layers", 0)):95            self.gguf_writer.add_nextn_predict_layers(nextn_layers)96 97    def prepare_metadata(self, vocab_only: bool):98        from_dir = self.fname_out.is_dir()99        super().prepare_metadata(vocab_only=vocab_only)100 101        if not self.mtp_only or not from_dir:102            return103 104        output_type: str = self.ftype.name.partition("_")[2]105        fname_default: str = gguf.naming_convention(106            self.metadata.name, self.metadata.basename, self.metadata.finetune,107            self.metadata.version, size_label=None, output_type=output_type, model_type=None)108        self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf"109 110    @classmethod111    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:112        name, gen = item113        if name.endswith(".expert_bias"):114            name += ".bias"115 116        if cls._main_layers is None:117            return super().filter_tensors((name, gen))118 119        m = re.match(r"model\.layers\.(\d+)\.", name)120        is_mtp = m is not None and int(m.group(1)) >= cls._main_layers121 122        if is_mtp and cls.no_mtp:123            return None124        if cls.mtp_only and not is_mtp and name not in (125            "model.word_embeddings.weight", "model.norm.weight", "lm_head.weight",126        ):127            return None128 129        return super().filter_tensors((name, gen))130 131    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:132        if name.endswith((".q_conv1d.weight", ".k_conv1d.weight", ".v_conv1d.weight")) and data_torch.ndim in (2, 3):133            d_inner = data_torch.shape[0]134            d_conv = data_torch.shape[-1]135            data_torch = data_torch.reshape(1, d_inner, 1, d_conv)136 137        if name.endswith(".A_log"):138            data_torch = torch.exp(data_torch).reshape(-1, 1)139 140        if name.endswith(".dt_bias"):141            name = name.rpartition(".dt_bias")[0] + ".dt_proj.bias"142 143        if name.endswith(".attention.f_proj.weight"):144            assert bid is not None145            if self.is_full_attention(bid):146                raise ValueError(f"unexpected f_proj on full-attention layer {bid}")147            name = self.format_tensor_name(gguf.MODEL_TENSOR.SSM_F_A, bid)148 149        if name.endswith(".attention.g_proj.weight"):150            assert bid is not None151            tensor = gguf.MODEL_TENSOR.ATTN_GATE if self.is_full_attention(bid) else gguf.MODEL_TENSOR.SSM_G_A152            name = self.format_tensor_name(tensor, bid)153 154        if ".mlp.experts." in name:155            n_experts = self.hparams["num_experts"]156            assert bid is not None157 158            if self._experts is None:159                self._experts = [{} for _ in range(self.block_count)]160 161            self._experts[bid][name] = data_torch162            if len(self._experts[bid]) >= n_experts * 3:163                for weight_name in ("down_proj", "gate_proj", "up_proj"):164                    tensors = []165                    for expert_id in range(n_experts):166                        expert_name = f"model.layers.{bid}.mlp.experts.{expert_id}.{weight_name}.weight"167                        tensors.append(self._experts[bid].pop(expert_name))168                    merged_name = f"model.layers.{bid}.mlp.experts.{weight_name}.weight"169                    yield from super().modify_tensors(torch.stack(tensors, dim=0), merged_name, bid)170            return171 172        if name.endswith(".attention.kv_b_proj.weight"):173            assert bid is not None174            n_head = self.hparams["num_attention_heads"]175            v_head_dim = self.hparams["v_head_dim"]176            qk_nope_head_dim = self.hparams["qk_nope_head_dim"]177            assert data_torch.shape[0] == n_head * (v_head_dim + qk_nope_head_dim)178            kv_b = data_torch.view(n_head, v_head_dim + qk_nope_head_dim, data_torch.shape[-1])179            k_b, v_b = torch.split(kv_b, [qk_nope_head_dim, v_head_dim], dim=1)180            name_k = self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_K_B, bid)181            name_v = self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_V_B, bid)182            yield from super().modify_tensors(k_b.transpose(1, 2), name_k, bid)183            yield from super().modify_tensors(v_b, name_v, bid)184            return185 186        yield from super().modify_tensors(data_torch, name, bid)187 188    def prepare_tensors(self):189        super().prepare_tensors()190        if self._experts is not None:191            experts = [name for layer in self._experts for name in layer]192            if experts:193                raise ValueError(f"Unprocessed experts: {experts}")194