Buckets:
| from __future__ import annotations | |
| from typing import Callable, Iterable, TYPE_CHECKING | |
| import torch | |
| if TYPE_CHECKING: | |
| from torch import Tensor | |
| from .base import ModelBase, gguf | |
| from .mistral import MistralModel, MistralMoeModel | |
| class MistralLarge4Model(MistralMoeModel): | |
| model_arch = gguf.MODEL_ARCH.MISTRAL_LARGE4 | |
| def __init__(self, *args, **kwargs): | |
| super().__init__(*args, **kwargs) | |
| n_layer = self.hparams["num_hidden_layers"] | |
| dsa = self.hparams["dsa"] | |
| def expand(pattern: str) -> list[str]: | |
| values = pattern.split(",") | |
| return [values[i % len(values)] for i in range(n_layer)] | |
| self.layer_has_rope = [v == "rope" for v in expand(self.hparams.get("nope") or "rope")] | |
| self.layer_indexer_full = [v == "full" for v in expand(dsa.get("sharing") or "full")] | |
| kv_sharing = self.hparams.get("kv_sharing") or {} | |
| if "source_layers_str" in kv_sharing: | |
| self.layer_kv_source = [int(v) for v in kv_sharing["source_layers_str"].split(",")] | |
| else: | |
| assert not kv_sharing, f"unsupported kv_sharing config: {kv_sharing}" | |
| self.layer_kv_source = list(range(n_layer)) | |
| assert len(self.layer_kv_source) == n_layer | |
| # note: the rope dim of the indexer must match the main attention | |
| assert dsa.get("rope_head_dim", -1) in (-1, self.hparams["qk_rope_head_dim"]) | |
| assert dsa.get("match_nope", False), "indexer must follow the NoPE pattern of the main attention" | |
| assert dsa.get("ssmax_scaling", False), "indexer without SSMax is not supported" | |
| assert self.hparams.get("q_lora_rank") is None | |
| assert not self.hparams["yarn"].get("apply_scale", True), "only YaRN without magnitude scaling is supported" | |
| def dequant_model(self): | |
| quant_config = self.hparams.get("quantization") | |
| if quant_config is not None: | |
| assert quant_config["qformat_weight"] == "fp8_e4m3" | |
| assert quant_config.get("qscheme_weight") == "BLOCK" | |
| self.hparams["quantization_config"] = { | |
| "quant_method": "fp8", | |
| "weight_block_size": [quant_config["block_size_out"], quant_config["block_size_in"]], | |
| } | |
| return super().dequant_model() | |
| def set_gguf_parameters(self): | |
| # skip the MistralMoeModel overrides, they apply the deepseek2 YaRN scaling | |
| super(MistralMoeModel, self).set_gguf_parameters() | |
| MistralModel.set_mistral_config(self.gguf_writer, self.hparams) | |
| dsa = self.hparams["dsa"] | |
| self.gguf_writer.add_indexer_head_count(dsa["n_heads"]) | |
| self.gguf_writer.add_indexer_key_length(dsa["head_dim"]) | |
| self.gguf_writer.add_indexer_top_k(dsa["topk"]) | |
| self.gguf_writer.add_indexer_types(self.layer_indexer_full) | |
| self.gguf_writer.add_layer_norm_eps(dsa.get("norm_eps", 1e-6)) | |
| self.gguf_writer.add_rope_pattern(self.layer_has_rope) | |
| self.gguf_writer.add_kv_source_map(self.layer_kv_source) | |
| self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SOFTMAX) | |
| self.gguf_writer.add_final_logit_softcapping(self.hparams["softmax_tanh_gating"]) | |
| def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: | |
| name, gen = item | |
| # exported separately with --dspark | |
| if name.startswith("mtp_"): | |
| return None | |
| # experts need HF names for the merging logic in DeepseekV2Model | |
| # note: bare "layers.N.experts" comes from the NVFP4 quantized_layers config | |
| if ".experts." in name or name.endswith(".experts"): | |
| name = "model." + name.replace(".experts", ".mlp.experts") | |
| name = name.replace(".w1.", ".gate_proj.").replace(".w2.", ".down_proj.").replace(".w3.", ".up_proj.") | |
| # skip the MistralMoeModel renaming, the qscale names are needed for block FP8 dequant | |
| return super(MistralMoeModel, cls).filter_tensors((name, gen)) | |
| def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: | |
| if name == "lerope.alpha": | |
| # LeRoPE scales inv_freq by exp(alpha), ggml divides theta by the freq factors | |
| yield (self.format_tensor_name(gguf.MODEL_TENSOR.ROPE_FREQS), torch.exp(-data_torch.float())) | |
| return | |
| if bid is not None and name.endswith((".ssmax_scaling.s", ".ssmax_scaling.b")): | |
| tensor = gguf.MODEL_TENSOR.INDEXER_SSMAX if ".indexer." in name else gguf.MODEL_TENSOR.ATTN_SSMAX | |
| suffix = ".weight" if name.endswith(".s") else ".bias" | |
| yield (self.format_tensor_name(tensor, bid, suffix), data_torch.float()) | |
| return | |
| if bid is not None and ".indexer.q_input_norm." in name: | |
| suffix = ".weight" if name.endswith(".weight") else ".bias" | |
| yield (self.format_tensor_name(gguf.MODEL_TENSOR.INDEXER_Q_NORM, bid, suffix), data_torch) | |
| return | |
| if bid is not None and ".indexer." in name: | |
| name = "model." + name.replace(".attention.", ".self_attn.") | |
| # NoPE layers have a larger non-rotary head, so wkv_b is split per layer | |
| if bid is not None and name.endswith(".wkv_b.weight"): | |
| n_head = self.hparams["num_attention_heads"] | |
| v_head_dim = self.hparams["v_head_dim"] | |
| qk_nope_head_dim = self.hparams["qk_nope_head_dim"] | |
| if not self.layer_has_rope[bid]: | |
| qk_nope_head_dim += self.hparams["qk_rope_head_dim"] | |
| 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 = torch.split(kv_b, [qk_nope_head_dim, v_head_dim], dim=1) | |
| k_b = k_b.transpose(1, 2) | |
| yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_K_B, bid), k_b) | |
| yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_V_B, bid), v_b) | |
| return | |
| yield from super().modify_tensors(data_torch, name, bid) | |
| class MistralLarge4DSparkModel(MistralLarge4Model): | |
| """Mistral Large 4 DSpark draft (mtp_dspark.*) as a DFLASH GGUF, export with --dspark""" | |
| model_arch = gguf.MODEL_ARCH.DFLASH | |
| _DSPARK_ROOT: dict[str, gguf.MODEL_TENSOR] = { | |
| "backbone.target_projection.weight": gguf.MODEL_TENSOR.FC, | |
| "backbone.target_norm.weight": gguf.MODEL_TENSOR.ENC_OUTPUT_NORM, | |
| "backbone.final_norm.weight": gguf.MODEL_TENSOR.OUTPUT_NORM, | |
| "backbone.mask_embedding": gguf.MODEL_TENSOR.DSPARK_MASK_EMBD, | |
| "autoregressive_head.embedding.weight": gguf.MODEL_TENSOR.DSPARK_MARKOV_W1, | |
| "autoregressive_head.output.weight": gguf.MODEL_TENSOR.DSPARK_MARKOV_W2, | |
| "autoregressive_head.joint_proj.weight": gguf.MODEL_TENSOR.DSPARK_JOINT_PROJ, | |
| "autoregressive_head.joint_proj.bias": gguf.MODEL_TENSOR.DSPARK_JOINT_PROJ, | |
| } | |
| def __init__(self, *args, **kwargs): | |
| super().__init__(*args, **kwargs) | |
| mtp = self.hparams.get("mtp") or {} | |
| dspark = mtp.get("dspark") | |
| if dspark is None: | |
| raise ValueError("Mistral Large 4 DSpark export needs the mtp.dspark config") | |
| assert dspark.get("head_type") == "rnn", f"unsupported DSpark head type: {dspark.get('head_type')}" | |
| assert dspark.get("draft_input") == "position_specific", f"unsupported DSpark draft input: {dspark.get('draft_input')}" | |
| windows = {int(v) for v in (mtp.get("ragged_attention") or "").split(",") if v.strip()} | |
| assert len(windows) == 1, f"DSpark draft layers must share one sliding window, got {mtp.get('ragged_attention')}" | |
| self._n_stage = int(mtp["n_layers"]) | |
| self._depth = int(mtp["depth"]) | |
| self._window = windows.pop() | |
| self._target_layers = [int(v) for v in mtp["target_activation_layers"].split(",")] | |
| self.layer_has_rope = [True] * self._n_stage | |
| def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: | |
| name, gen = item | |
| if not name.startswith("mtp_dspark."): | |
| return None | |
| name = name.removeprefix("mtp_dspark.") | |
| if name.startswith("backbone.layers."): | |
| return name.removeprefix("backbone."), gen | |
| return f"dspark.{name}", gen | |
| def set_vocab(self): | |
| super().set_vocab() | |
| # block uses mask_embd, mask token id is not used | |
| self.gguf_writer.add_mask_token_id(0) | |
| def set_gguf_parameters(self): | |
| hp = self.hparams | |
| w = self.gguf_writer | |
| w.add_block_count(self._n_stage) | |
| w.add_context_length(hp["max_position_embeddings"]) | |
| w.add_embedding_length(hp["hidden_size"]) | |
| w.add_feed_forward_length(hp["intermediate_size"]) | |
| w.add_head_count(hp["num_attention_heads"]) | |
| w.add_head_count_kv(1) | |
| w.add_key_length(hp["kv_lora_rank"] + hp["qk_rope_head_dim"]) | |
| w.add_value_length(hp["kv_lora_rank"]) | |
| w.add_key_length_mla(hp["qk_nope_head_dim"] + hp["qk_rope_head_dim"]) | |
| w.add_value_length_mla(hp["v_head_dim"]) | |
| w.add_kv_lora_rank(hp["kv_lora_rank"]) | |
| w.add_rope_dimension_count(hp["qk_rope_head_dim"]) | |
| w.add_rope_freq_base(float(hp["rope_theta"])) | |
| w.add_layer_norm_rms_eps(hp["rms_norm_eps"]) | |
| MistralModel.set_mistral_config(w, hp) | |
| w.add_sliding_window(self._window) | |
| w.add_final_logit_softcapping(hp["softmax_tanh_gating"]) | |
| w.add_block_size(self._depth) | |
| # draft reads layer output = input of next layer | |
| w.add_target_layers([layer + 1 for layer in self._target_layers]) | |
| w.add_sample_from_anchor(True) | |
| w.add_has_confidence_head(False) | |
| w.add_file_type(self.ftype) | |
| def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: | |
| if name.startswith("dspark."): | |
| key = name.removeprefix("dspark.") | |
| if key not in self._DSPARK_ROOT: | |
| raise ValueError(f"Unexpected DSpark tensor {name!r}") | |
| suffix = ".bias" if key.endswith(".bias") else ".weight" | |
| yield (self.format_tensor_name(self._DSPARK_ROOT[key], None, suffix), data_torch) | |
| return | |
| yield from super().modify_tensors(data_torch, name, bid) | |
Xet Storage Details
- Size:
- 10.5 kB
- Xet hash:
- 5f2097bbc56c7ff1aeac1ea62345f2b87a1c167940f81add8c0d39fcb29aeb09
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.