Felipe97/llama-cpp-compiled

1from __future__ import annotations2 3import math4import re5 6import torch7 8from typing import TYPE_CHECKING, Any, Callable, Iterable9 10if TYPE_CHECKING:11    from torch import Tensor12 13from .base import MmprojModel, ModelBase, gguf14 15from .deepseek import DeepseekV2Model16 17 18@ModelBase.register("Dots3NoteForCausalLM", "Dots3NoteForConditionalGeneration", "Dots3NoteTextForCausalLM")19class Dots3NoteModel(DeepseekV2Model):20    model_arch = gguf.MODEL_ARCH.DOTS3NOTE21    skip_mtp = False22    supports_mtp_export = True23 24    # trunk layer count, stashed before indexing for filter_tensors (mirrors DeepseekV32Model)25    _n_main_layers: int | None = None26 27    def index_tensors(self, remote_hf_model_id: str | None = None):28        type(self)._n_main_layers = self.hparams["num_hidden_layers"]29        return super().index_tensors(remote_hf_model_id=remote_hf_model_id)30 31    def __init__(self, *args, **kwargs):32        super().__init__(*args, **kwargs)33 34        hparams = self.hparams35 36        # config file doesn't specify MTP block, detect it from model weight37        self.n_nextn = 1 if "model.mtp.embed_tokens.weight" in self.model_tensors else 038        if self.n_nextn:39            self.block_count += self.n_nextn40            self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)41 42        self.layer_types = hparams["layer_types"]43        if len(self.layer_types) < hparams["num_hidden_layers"]:44            raise ValueError("layer_types is shorter than num_hidden_layers")45 46        if hparams.get("use_dsa", True) is not True:47            raise ValueError("dots3-note conversion requires use_dsa=true")48        if hparams.get("normalization", "RMSNorm") != "RMSNorm" or hparams.get("final_norm", "RMSNorm") != "RMSNorm":49            raise ValueError("dots3-note conversion only supports RMSNorm")50        if hparams.get("k_rope_only_layernorm", True) is not True:51            raise ValueError("dots3-note conversion requires k_rope_only_layernorm=true")52        if hparams.get("topk_method", "noaux_tc") != "noaux_tc" or hparams.get("scoring_func") != "sigmoid":53            raise ValueError("dots3-note conversion only supports noaux_tc/sigmoid expert gating")54        if hparams.get("n_group", 1) != 1 or hparams.get("topk_group", 1) != 1:55            raise ValueError("dots3-note conversion does not support grouped expert routing")56        if hparams.get("use_dynamic_rsf", False) or hparams.get("moe_gating_fp32", False):57            raise ValueError("dots3-note conversion does not support use_dynamic_rsf/moe_gating_fp32")58        for key in ("attention_gate_type", "swa_attention_gate_type"):59            if hparams.get(key, "headwise") != "headwise":60                raise ValueError(f"dots3-note conversion only supports headwise attention gate, got {key}={hparams.get(key)!r}")61        if hparams["swa_qk_nope_head_dim"] + hparams["swa_qk_rope_head_dim"] != hparams.get("swa_head_dim", 256):62            raise ValueError("swa_head_dim must equal swa_qk_nope_head_dim + swa_qk_rope_head_dim")63        if hparams["swa_qk_rope_head_dim"] != hparams["qk_rope_head_dim"]:64            # both layer kinds share a single rope_dimension_count65            raise ValueError("swa_qk_rope_head_dim must match qk_rope_head_dim")66 67        self.apply_lora_rescale = hparams.get("apply_mla_qkv_lora_rescale", False)68 69    def _is_swa_layer(self, bid: int) -> bool:70        if bid >= self.hparams["num_hidden_layers"]:71            # note: the NextN/MTP block uses the sliding-attention MLA72            return True73        return self.layer_types[bid] == "sliding_attention"74 75    def set_vocab(self):76        from transformers import AutoTokenizer77        tokenizer = AutoTokenizer.from_pretrained(self.dir_model)78        special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True)79        tokens, toktypes, tokpre = self.get_vocab_base()80        self.gguf_writer.add_tokenizer_model("gpt2")81        self.gguf_writer.add_tokenizer_pre(tokpre)82        self.gguf_writer.add_token_list(tokens)83        self.gguf_writer.add_token_types(toktypes)84        special_vocab._set_special_token("eot", tokenizer.get_added_vocab()["<|endofassistant|>"])  # ty: ignore[unresolved-attribute]85        special_vocab.add_to_gguf(self.gguf_writer)86 87    @classmethod88    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:89        if (titem := super().filter_tensors(item)) is None:90            return None91        name, gen = titem92        if name.startswith(("vision_encoder.", "audio_encoder.")):93            return None94 95        assert cls._n_main_layers is not None96        is_mtp = name.startswith("model.mtp.") or \97            ((m := re.match(r"model\.layers\.(\d+)\.", name)) is not None and int(m.group(1)) >= cls._n_main_layers)98 99        # --no-mtp: drop the NextN/MTP block; --mtp: keep only that block plus the shared embeddings/norm/lm_head100        if is_mtp and cls.no_mtp:101            return None102        if cls.mtp_only and not is_mtp and name not in (103            "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight",104        ):105            return None106 107        return name, gen108 109    def set_gguf_parameters(self):110        hparams = self.hparams111 112        # head_count is a per-layer array because the two layer kinds have different head counts113        n_layer = hparams["num_hidden_layers"]114        hparams["num_attention_heads"] = [115            hparams["swa_num_attention_heads"] if self._is_swa_layer(il) else hparams["num_attention_heads"]116            for il in range(self.block_count)117        ]118 119        # prevent the base class from emitting key/value_length from the unused head_dim120        hparams.pop("head_dim", None)121 122        super().set_gguf_parameters()123 124        # MLA geometry of the sliding-window layers (rope.freq_base_swa is emitted by the base class)125        swa_kv_lora_rank = hparams["swa_kv_lora_rank"]126        self.gguf_writer.add_sliding_window(hparams["sliding_window_size"])127        self.gguf_writer.add_sliding_window_pattern([self._is_swa_layer(il) for il in range(n_layer)])128        self.gguf_writer.add_kv_lora_rank_swa(swa_kv_lora_rank)129        self.gguf_writer.add_key_length_swa(swa_kv_lora_rank + hparams["swa_qk_rope_head_dim"])130        self.gguf_writer.add_value_length_swa(swa_kv_lora_rank)131        self.gguf_writer.add_key_length_mla_swa(hparams["swa_qk_nope_head_dim"] + hparams["swa_qk_rope_head_dim"])132        self.gguf_writer.add_value_length_mla_swa(hparams["swa_v_head_dim"])133        if hparams["swa_q_lora_rank"] != hparams["q_lora_rank"]:134            raise ValueError("dots3-note conversion assumes a shared q_lora_rank for both layer kinds")135 136        if self.n_nextn:137            self.gguf_writer.add_nextn_predict_layers(self.n_nextn)138 139        # DSA indexer (full-attention layers only)140        self.gguf_writer.add_indexer_head_count(hparams["index_n_heads"])141        self.gguf_writer.add_indexer_key_length(hparams["index_head_dim"])142        self.gguf_writer.add_indexer_top_k(hparams["index_topk"])143        self.gguf_writer.add_indexer_types([not self._is_swa_layer(il) for il in range(n_layer)])144 145    def prepare_metadata(self, vocab_only: bool):146        from_dir = self.fname_out.is_dir()147        super().prepare_metadata(vocab_only=vocab_only)148 149        if not self.mtp_only or not from_dir:150            return151 152        output_type: str = self.ftype.name.partition("_")[2]153        fname_default: str = gguf.naming_convention(154            self.metadata.name, self.metadata.basename, self.metadata.finetune,155            self.metadata.version, size_label=None, output_type=output_type, model_type=None)156        self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf"157 158    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:159        # move the MTP token embedding into the NextN block so the standard nextn mapping picks it up160        if name == "model.mtp.embed_tokens.weight":161            name = f"model.layers.{self.hparams['num_hidden_layers']}.embed_tokens.weight"162            bid = self.hparams["num_hidden_layers"]163 164        # fold the activation rescale sqrt(n_embd/lora_rank) into the preceding RMSNorm weight165        # this also covers the indexer wq_b, which reads the same rescaled q_lora activation166        if self.apply_lora_rescale and bid is not None:167            if name.endswith("q_a_layernorm.weight"):168                data_torch = data_torch * math.sqrt(self.hparams["hidden_size"] / self.hparams["q_lora_rank"])169            elif name.endswith("kv_a_layernorm.weight"):170                rank = self.hparams["swa_kv_lora_rank"] if self._is_swa_layer(bid) else self.hparams["kv_lora_rank"]171                data_torch = data_torch * math.sqrt(self.hparams["hidden_size"] / rank)172 173        # MLA absorption: split kv_b_proj into k_b (transposed) and v_b, per-layer-kind geometry174        if name.endswith("kv_b_proj.weight"):175            assert bid is not None176            if self._is_swa_layer(bid):177                n_head = self.hparams["swa_num_attention_heads"]178                qk_nope_head_dim = self.hparams["swa_qk_nope_head_dim"]179                v_head_dim = self.hparams["swa_v_head_dim"]180            else:181                n_head = self.hparams["num_attention_heads"]182                qk_nope_head_dim = self.hparams["qk_nope_head_dim"]183                v_head_dim = self.hparams["v_head_dim"]184            if isinstance(n_head, list):  # set_gguf_parameters turns this into a per-layer array185                n_head = n_head[bid]186 187            assert data_torch.shape[0] == n_head * (qk_nope_head_dim + v_head_dim)188 189            kv_b = data_torch.view(n_head, qk_nope_head_dim + v_head_dim, data_torch.shape[-1])190            k_b, v_b = kv_b.split([qk_nope_head_dim, v_head_dim], dim=1)191            k_b = k_b.transpose(1, 2)192 193            yield from ModelBase.modify_tensors(self, k_b, name.replace("kv_b_proj", "k_b_proj"), bid)194            yield from ModelBase.modify_tensors(self, v_b, name.replace("kv_b_proj", "v_b_proj"), bid)195            return196 197        yield from super().modify_tensors(data_torch, name, bid)198 199 200@ModelBase.register("Dots3NoteForCausalLM", "Dots3NoteForConditionalGeneration")201class Dots3NoteMmprojModel(MmprojModel):202    has_vision_encoder = True203    has_audio_encoder = True204 205    def __init__(self, *args, **kwargs):206        super().__init__(*args, **kwargs)207        assert self.hparams_vision is not None208        assert self.hparams_audio is not None209 210        # preprocessor_config.json nests the image params under vision_config211        self.preprocessor_config = {**self.preprocessor_config, **self.preprocessor_config.get("vision_config", {})}212 213        vis = self.hparams_vision214        # in this config, hidden_size is the adapter output width; embed_dim is the tower width215        vis["hidden_size"] = vis["embed_dim"]216        vis["image_size"] = 0  # dynamic resolution217        self.pyramid = [max(0, n) for n in vis["pyramid_num_routed"]]218 219        if vis.get("adapter_type") != "patch_merger" or not vis.get("pre_pixel_shuffle"):220            raise ValueError("dots3-note vision conversion requires adapter_type=patch_merger and pre_pixel_shuffle")221        if vis.get("router_scoring_func", "sigmoid") != "sigmoid" or vis.get("router_scale", 1.0) != 1.0:222            raise ValueError("dots3-note vision conversion only supports sigmoid routing with router_scale=1.0")223        if vis.get("temporal_patch_size", 1) != 1 or vis.get("use_bias") or not vis.get("use_qk_norm"):224            raise ValueError("unsupported dots3-note vision config variant")225 226        aud = self.hparams_audio227        if not aud.get("use_conv2d_stem") or not aud.get("use_rope") or not aud.get("use_rms_norm") or aud.get("use_causal"):228            raise ValueError("unsupported dots3-note audio config variant")229        if aud["whisper_config"].get("activation_function") != "swiglu":230            raise ValueError("dots3-note audio conversion requires the swiglu activation")231        if aud.get("merge_factor", 1) != 1 or aud.get("chunk_seconds") != 60:232            raise ValueError("unsupported dots3-note audio chunking config")233        # the graph hard-codes these rope parameters234        rope = aud.get("rope_parameters", {})235        if rope.get("partial_rotary_factor") != 0.5 or rope.get("rope_theta") != 10000.0:236            raise ValueError("unsupported dots3-note audio rope config")237 238    def get_audio_config(self) -> dict[str, Any] | None:239        cfg = self.global_config.get("audio_config")240        if cfg is not None:241            # aliases so MmprojModel.find_aparam() / n_block_keys can resolve them242            whisper = cfg["whisper_config"]243            cfg["hidden_size"] = whisper["d_model"]244            cfg["intermediate_size"] = whisper["encoder_ffn_dim"]245            cfg["num_attention_heads"] = whisper["encoder_attention_heads"]246            cfg["num_hidden_layers"] = whisper["encoder_layers"]247        return cfg248 249    def set_gguf_parameters(self):250        super().set_gguf_parameters()251        assert self.hparams_vision is not None252        assert self.hparams_audio is not None253 254        self.gguf_writer.add_clip_vision_projector_type(gguf.VisionProjectorType.DOTS3NOTE_V)255        self.gguf_writer.add_vision_use_silu(True)256        self.gguf_writer.add_vision_attention_layernorm_eps(self.hparams_vision["rms_norm_eps"])257        self.gguf_writer.add_vision_spatial_merge_size(self.hparams_vision["spatial_merge_size"])258        self.gguf_writer.add_vision_min_pixels(self.preprocessor_config["min_pixels"])259        self.gguf_writer.add_vision_max_pixels(self.preprocessor_config["max_pixels"])260        # pyramid MoE: per-block routed expert count, 0 = dense block261        self.gguf_writer.add_vision_expert_count_per_layer(self.pyramid)262        self.gguf_writer.add_vision_expert_used_count(int(self.hparams_vision["capacity_factor"]))263 264        self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.DOTS3NOTE_A)265        self.gguf_writer.add_audio_num_mel_bins(self.hparams_audio["whisper_config"]["num_mel_bins"])266        self.gguf_writer.add_audio_attention_layernorm_eps(1e-6)  # Dots3NoteAudioRMSNorm default267 268    @classmethod269    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:270        name, _ = item271        if not name.startswith(("vision_encoder.", "audio_encoder.")):272            return None273        return super().filter_tensors(item)274 275    _vis_experts: dict[int, dict[str, Tensor]] | None = None276 277    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:278        # router params have no .weight suffix in the checkpoint, but gguf tools expect one279        if name.endswith((".gate_weight", ".router_bias")):280            name += ".weight"281 282        # audio fc1 fuses gate and up for swiglu; split it283        if ".speech_encoder.layers." in name and ".fc1." in name:284            gate, up = data_torch.chunk(2, dim=0)285            yield from super().modify_tensors(gate, name.replace(".fc1.", ".fc1_gate."), bid)286            yield from super().modify_tensors(up, name.replace(".fc1.", ".fc1_up."), bid)287            return288 289        # vision MoE: stack per-expert weights into a single 3D tensor per block290        if ".mlp.experts." in name:291            assert bid is not None292            n_expert = self.pyramid[bid]293            if self._vis_experts is None:294                self._vis_experts = {}295            buf = self._vis_experts.setdefault(bid, {})296            buf[name] = data_torch297 298            if len(buf) >= n_expert * 3:299                for w_name in ("fc1", "fc2", "fc3"):300                    datas: list[Tensor] = []301                    for xid in range(n_expert):302                        ename = f"vision_encoder.blocks.{bid}.mlp.experts.{xid}.{w_name}.weight"303                        datas.append(buf.pop(ename))304                    merged = torch.stack(datas, dim=0)305                    yield from super().modify_tensors(merged, f"vision_encoder.blocks.{bid}.mlp.experts.{w_name}.weight", bid)306            return307 308        yield from super().modify_tensors(data_torch, name, bid)309 310    def prepare_tensors(self):311        super().prepare_tensors()312        if self._vis_experts is not None:313            leftover = [k for d in self._vis_experts.values() for k in d.keys()]314            if leftover:315                raise ValueError(f"unprocessed vision experts: {leftover}")316 317    def tensor_force_quant(self, name, new_name, bid, n_dims):318        # FP32 routing is load-bearing for the vision MoE (near-tied expert scores)319        if ".ffn_gate_inp." in new_name or ".exp_probs_b." in new_name:320            return gguf.GGMLQuantizationType.F32321        if ".conv2d" in new_name or "a.conv_out" in new_name:322            return gguf.GGMLQuantizationType.F32323        return super().tensor_force_quant(name, new_name, bid, n_dims)324