diff --git a/conversion/__init__.py b/conversion/__init__.py index 695289b73a..263bb2c347 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -60,6 +60,9 @@ TEXT_MODEL_MAP: dict[str, str] = { "DistilBertForSequenceClassification": "bert", "DistilBertModel": "bert", "Dots1ForCausalLM": "dots1", + "Dots3NoteForCausalLM": "dots3", + "Dots3NoteForConditionalGeneration": "dots3", + "Dots3NoteTextForCausalLM": "dots3", "DotsOCRForCausalLM": "qwen", "DreamModel": "dream", "Ernie4_5ForCausalLM": "ernie", diff --git a/conversion/dots3.py b/conversion/dots3.py new file mode 100644 index 0000000000..93362537ee --- /dev/null +++ b/conversion/dots3.py @@ -0,0 +1,158 @@ +from __future__ import annotations + +import math + +from typing import TYPE_CHECKING, Callable, Iterable + +if TYPE_CHECKING: + from torch import Tensor + +from .base import ModelBase, gguf + +from .deepseek import DeepseekV2Model + + +@ModelBase.register("Dots3NoteForCausalLM", "Dots3NoteForConditionalGeneration", "Dots3NoteTextForCausalLM") +class Dots3NoteModel(DeepseekV2Model): + model_arch = gguf.MODEL_ARCH.DOTS3NOTE + skip_mtp = False + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + hparams = self.hparams + + # config file doesn't specify MTP block, detect it from model weight + self.n_nextn = 1 if "model.mtp.embed_tokens.weight" in self.model_tensors else 0 + if self.n_nextn: + self.block_count += self.n_nextn + self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count) + + self.layer_types = hparams["layer_types"] + if len(self.layer_types) < hparams["num_hidden_layers"]: + raise ValueError("layer_types is shorter than num_hidden_layers") + + if hparams.get("use_dsa", True) is not True: + raise ValueError("dots3-note conversion requires use_dsa=true") + if hparams.get("normalization", "RMSNorm") != "RMSNorm" or hparams.get("final_norm", "RMSNorm") != "RMSNorm": + raise ValueError("dots3-note conversion only supports RMSNorm") + if hparams.get("k_rope_only_layernorm", True) is not True: + raise ValueError("dots3-note conversion requires k_rope_only_layernorm=true") + if hparams.get("topk_method", "noaux_tc") != "noaux_tc" or hparams.get("scoring_func") != "sigmoid": + raise ValueError("dots3-note conversion only supports noaux_tc/sigmoid expert gating") + if hparams.get("n_group", 1) != 1 or hparams.get("topk_group", 1) != 1: + raise ValueError("dots3-note conversion does not support grouped expert routing") + if hparams.get("use_dynamic_rsf", False) or hparams.get("moe_gating_fp32", False): + raise ValueError("dots3-note conversion does not support use_dynamic_rsf/moe_gating_fp32") + for key in ("attention_gate_type", "swa_attention_gate_type"): + if hparams.get(key, "headwise") != "headwise": + raise ValueError(f"dots3-note conversion only supports headwise attention gate, got {key}={hparams.get(key)!r}") + if hparams["swa_qk_nope_head_dim"] + hparams["swa_qk_rope_head_dim"] != hparams.get("swa_head_dim", 256): + raise ValueError("swa_head_dim must equal swa_qk_nope_head_dim + swa_qk_rope_head_dim") + if hparams["swa_qk_rope_head_dim"] != hparams["qk_rope_head_dim"]: + # both layer kinds share a single rope_dimension_count + raise ValueError("swa_qk_rope_head_dim must match qk_rope_head_dim") + + self.apply_lora_rescale = hparams.get("apply_mla_qkv_lora_rescale", False) + + def _is_swa_layer(self, bid: int) -> bool: + if bid >= self.hparams["num_hidden_layers"]: + # note: the NextN/MTP block uses the sliding-attention MLA + return True + return self.layer_types[bid] == "sliding_attention" + + def set_vocab(self): + from transformers import AutoTokenizer + tokenizer = AutoTokenizer.from_pretrained(self.dir_model) + special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True) + tokens, toktypes, tokpre = self.get_vocab_base() + self.gguf_writer.add_tokenizer_model("gpt2") + self.gguf_writer.add_tokenizer_pre(tokpre) + self.gguf_writer.add_token_list(tokens) + self.gguf_writer.add_token_types(toktypes) + special_vocab._set_special_token("eot", tokenizer.get_added_vocab()["<|endofassistant|>"]) # ty: ignore[unresolved-attribute] + special_vocab.add_to_gguf(self.gguf_writer) + + @classmethod + def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: + name, _ = item + if name.startswith(("vision_encoder.", "audio_encoder.")): + return None + return super().filter_tensors(item) + + def set_gguf_parameters(self): + hparams = self.hparams + + # head_count is a per-layer array because the two layer kinds have different head counts + n_layer = hparams["num_hidden_layers"] + hparams["num_attention_heads"] = [ + hparams["swa_num_attention_heads"] if self._is_swa_layer(il) else hparams["num_attention_heads"] + for il in range(self.block_count) + ] + + # prevent the base class from emitting key/value_length from the unused head_dim + hparams.pop("head_dim", None) + + super().set_gguf_parameters() + + # MLA geometry of the sliding-window layers (rope.freq_base_swa is emitted by the base class) + swa_kv_lora_rank = hparams["swa_kv_lora_rank"] + self.gguf_writer.add_sliding_window(hparams["sliding_window_size"]) + self.gguf_writer.add_sliding_window_pattern([self._is_swa_layer(il) for il in range(n_layer)]) + self.gguf_writer.add_kv_lora_rank_swa(swa_kv_lora_rank) + self.gguf_writer.add_key_length_swa(swa_kv_lora_rank + hparams["swa_qk_rope_head_dim"]) + self.gguf_writer.add_value_length_swa(swa_kv_lora_rank) + self.gguf_writer.add_key_length_mla_swa(hparams["swa_qk_nope_head_dim"] + hparams["swa_qk_rope_head_dim"]) + self.gguf_writer.add_value_length_mla_swa(hparams["swa_v_head_dim"]) + if hparams["swa_q_lora_rank"] != hparams["q_lora_rank"]: + raise ValueError("dots3-note conversion assumes a shared q_lora_rank for both layer kinds") + + if self.n_nextn: + self.gguf_writer.add_nextn_predict_layers(self.n_nextn) + + # DSA indexer (full-attention layers only) + self.gguf_writer.add_indexer_head_count(hparams["index_n_heads"]) + self.gguf_writer.add_indexer_key_length(hparams["index_head_dim"]) + self.gguf_writer.add_indexer_top_k(hparams["index_topk"]) + self.gguf_writer.add_indexer_types([not self._is_swa_layer(il) for il in range(n_layer)]) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + # move the MTP token embedding into the NextN block so the standard nextn mapping picks it up + if name == "model.mtp.embed_tokens.weight": + name = f"model.layers.{self.hparams['num_hidden_layers']}.embed_tokens.weight" + bid = self.hparams["num_hidden_layers"] + + # fold the activation rescale sqrt(n_embd/lora_rank) into the preceding RMSNorm weight + # this also covers the indexer wq_b, which reads the same rescaled q_lora activation + if self.apply_lora_rescale and bid is not None: + if name.endswith("q_a_layernorm.weight"): + data_torch = data_torch * math.sqrt(self.hparams["hidden_size"] / self.hparams["q_lora_rank"]) + elif name.endswith("kv_a_layernorm.weight"): + rank = self.hparams["swa_kv_lora_rank"] if self._is_swa_layer(bid) else self.hparams["kv_lora_rank"] + data_torch = data_torch * math.sqrt(self.hparams["hidden_size"] / rank) + + # MLA absorption: split kv_b_proj into k_b (transposed) and v_b, per-layer-kind geometry + if name.endswith("kv_b_proj.weight"): + assert bid is not None + if self._is_swa_layer(bid): + n_head = self.hparams["swa_num_attention_heads"] + qk_nope_head_dim = self.hparams["swa_qk_nope_head_dim"] + v_head_dim = self.hparams["swa_v_head_dim"] + else: + n_head = self.hparams["num_attention_heads"] + qk_nope_head_dim = self.hparams["qk_nope_head_dim"] + v_head_dim = self.hparams["v_head_dim"] + if isinstance(n_head, list): # set_gguf_parameters turns this into a per-layer array + n_head = n_head[bid] + + assert data_torch.shape[0] == n_head * (qk_nope_head_dim + v_head_dim) + + kv_b = data_torch.view(n_head, qk_nope_head_dim + v_head_dim, data_torch.shape[-1]) + k_b, v_b = kv_b.split([qk_nope_head_dim, v_head_dim], dim=1) + k_b = k_b.transpose(1, 2) + + yield from ModelBase.modify_tensors(self, k_b, name.replace("kv_b_proj", "k_b_proj"), bid) + yield from ModelBase.modify_tensors(self, v_b, name.replace("kv_b_proj", "v_b_proj"), bid) + return + + yield from super().modify_tensors(data_torch, name, bid) diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 98c4fa1691..3bce0357df 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -203,6 +203,9 @@ class Keys: VALUE_LENGTH_MLA = "{arch}.attention.value_length_mla" KEY_LENGTH_SWA = "{arch}.attention.key_length_swa" VALUE_LENGTH_SWA = "{arch}.attention.value_length_swa" + KEY_LENGTH_MLA_SWA = "{arch}.attention.key_length_mla_swa" + VALUE_LENGTH_MLA_SWA = "{arch}.attention.value_length_mla_swa" + KV_LORA_RANK_SWA = "{arch}.attention.kv_lora_rank_swa" SHARED_KV_LAYERS = "{arch}.attention.shared_kv_layers" SLIDING_WINDOW_PATTERN = "{arch}.attention.sliding_window_pattern" TEMPERATURE_SCALE = "{arch}.attention.temperature_scale" @@ -544,6 +547,7 @@ class MODEL_ARCH(IntEnum): BAILINGMOE = auto() BAILINGMOE2 = auto() DOTS1 = auto() + DOTS3NOTE = auto() ARCEE = auto() AFMOE = auto() LAGUNA = auto() @@ -1250,6 +1254,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = { MODEL_ARCH.BAILINGMOE: "bailingmoe", MODEL_ARCH.BAILINGMOE2: "bailingmoe2", MODEL_ARCH.DOTS1: "dots1", + MODEL_ARCH.DOTS3NOTE: "dots3note", MODEL_ARCH.ARCEE: "arcee", MODEL_ARCH.AFMOE: "afmoe", MODEL_ARCH.LAGUNA: "laguna", @@ -4231,6 +4236,44 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.FFN_UP_EXP, MODEL_TENSOR.FFN_UP_SHEXP, ], + MODEL_ARCH.DOTS3NOTE: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q_A, + MODEL_TENSOR.ATTN_Q_B, + MODEL_TENSOR.ATTN_KV_A_MQA, + MODEL_TENSOR.ATTN_K_B, + MODEL_TENSOR.ATTN_V_B, + MODEL_TENSOR.ATTN_Q_A_NORM, + MODEL_TENSOR.ATTN_KV_A_NORM, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_GATE, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.INDEXER_K_NORM, + MODEL_TENSOR.INDEXER_PROJ, + MODEL_TENSOR.INDEXER_ATTN_K, + MODEL_TENSOR.INDEXER_ATTN_Q_B, + # NextN/MTP tensors - preserved but unused + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_EMBED_TOKENS, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, + ], MODEL_ARCH.ARCEE: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 05f86396dc..3260d9d78e 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -785,6 +785,15 @@ class GGUFWriter: def add_key_length_swa(self, length: int) -> None: self.add_uint32(Keys.Attention.KEY_LENGTH_SWA.format(arch=self.arch), length) + def add_key_length_mla_swa(self, length: int) -> None: + self.add_uint32(Keys.Attention.KEY_LENGTH_MLA_SWA.format(arch=self.arch), length) + + def add_value_length_mla_swa(self, length: int) -> None: + self.add_uint32(Keys.Attention.VALUE_LENGTH_MLA_SWA.format(arch=self.arch), length) + + def add_kv_lora_rank_swa(self, length: int) -> None: + self.add_uint32(Keys.Attention.KV_LORA_RANK_SWA.format(arch=self.arch), length) + def add_value_length_swa(self, length: int) -> None: self.add_uint32(Keys.Attention.VALUE_LENGTH_SWA.format(arch=self.arch), length) diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index 79d270ab8f..ebe7621614 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -715,6 +715,7 @@ class TensorNameMap: "model.layers.layers.{bid}.mixer.k", # plamo2 "model.layers.layers.{bid}.mixer.k_norm", # plamo3 "layers.{bid}.self_attn.k_norm", # qwen3-embedding + "model.layers.{bid}.self_attn.k_rope_only_layernorm", # dots3note "model.layers.{bid}.attention.key_layernorm", # apertus ),