|
| 1 | +from __future__ import annotations |
| 2 | + |
| 3 | +import json |
| 4 | +from typing import Any, Iterable, TYPE_CHECKING |
| 5 | + |
| 6 | +import torch |
| 7 | + |
| 8 | +if TYPE_CHECKING: |
| 9 | + from torch import Tensor |
| 10 | + |
| 11 | +from .base import MmprojModel, ModelBase, TextModel, gguf |
| 12 | + |
| 13 | + |
| 14 | +def _unpermute_for_rope(tensor: "Tensor", n_heads: int) -> "Tensor": |
| 15 | + """Invert transformers' `_permute_for_rope`: HF stores Q/K in rotate_half layout, |
| 16 | + llama.cpp consumes the interleaved (NORM) layout.""" |
| 17 | + if tensor.ndim == 2: |
| 18 | + dim1, dim2 = tensor.shape |
| 19 | + return tensor.view(n_heads, 2, dim1 // n_heads // 2, dim2).transpose(1, 2).reshape(dim1, dim2) |
| 20 | + if tensor.ndim == 1: |
| 21 | + (dim1,) = tensor.shape |
| 22 | + return tensor.view(n_heads, 2, dim1 // n_heads // 2).transpose(1, 2).reshape(dim1) |
| 23 | + raise ValueError(f"_unpermute_for_rope: unexpected shape {tuple(tensor.shape)}") |
| 24 | + |
| 25 | + |
| 26 | +@ModelBase.register("MuseGlimmerForConditionalGeneration") |
| 27 | +class MuseGlimmerModel(TextModel): |
| 28 | + model_arch = gguf.MODEL_ARCH.MUSE_GLIMMER |
| 29 | + |
| 30 | + def norm_shift(self, name: str) -> float: |
| 31 | + # All four layer norms use 1, the final norm uses 0. |
| 32 | + return 1.0 if name.endswith("layernorm.weight") else 0.0 |
| 33 | + |
| 34 | + def set_vocab(self): |
| 35 | + self._set_vocab_gpt2() |
| 36 | + |
| 37 | + from transformers import AutoTokenizer |
| 38 | + tok = AutoTokenizer.from_pretrained(self.dir_model) |
| 39 | + eot_id = tok.convert_tokens_to_ids("<|eot|>") |
| 40 | + if isinstance(eot_id, int) and eot_id >= 0: |
| 41 | + self.gguf_writer.add_eot_token_id(eot_id) |
| 42 | + |
| 43 | + def set_gguf_parameters(self): |
| 44 | + super().set_gguf_parameters() |
| 45 | + hparams = self.hparams |
| 46 | + |
| 47 | + self.gguf_writer.add_final_logit_softcapping(hparams["final_logit_softcapping"]) |
| 48 | + self.gguf_writer.add_logit_scale(hparams["output_multiplier"]) |
| 49 | + self.gguf_writer.add_sliding_window(hparams["sliding_window"]) |
| 50 | + self.gguf_writer.add_sliding_window_pattern([t == "sliding_attention" for t in hparams["layer_types"]]) |
| 51 | + |
| 52 | + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: |
| 53 | + shift = self.norm_shift(name) |
| 54 | + if shift != 0.0: |
| 55 | + data_torch = data_torch + shift |
| 56 | + |
| 57 | + # Invert transformers' `_permute_for_rope` on Q/K, we keep ggml's NORM (interleaved) rope |
| 58 | + if ".self_attn.q_proj." in name: |
| 59 | + data_torch = _unpermute_for_rope(data_torch, int(self.hparams["num_attention_heads"])) |
| 60 | + elif ".self_attn.k_proj." in name: |
| 61 | + data_torch = _unpermute_for_rope(data_torch, int(self.hparams["num_key_value_heads"])) |
| 62 | + |
| 63 | + # Synthesize QK-norm weights to absorb qk_scale_factor. |
| 64 | + # MuseGlimmer implementation: scaleless RMSNorm followed by qk_scale_factor.. |
| 65 | + if bid is not None and name.endswith(f"model.layers.{bid}.self_attn.q_proj.weight"): |
| 66 | + head_dim = self.hparams["head_dim"] |
| 67 | + q_scale = float(self.hparams["qk_scale_factor"]) |
| 68 | + yield ( |
| 69 | + self.map_tensor_name(f"model.layers.{bid}.self_attn.q_norm.weight"), |
| 70 | + torch.full((head_dim,), q_scale, dtype=torch.float32), |
| 71 | + ) |
| 72 | + yield ( |
| 73 | + self.map_tensor_name(f"model.layers.{bid}.self_attn.k_norm.weight"), |
| 74 | + torch.ones((head_dim,), dtype=torch.float32), |
| 75 | + ) |
| 76 | + |
| 77 | + yield from super().modify_tensors(data_torch, name, bid) |
| 78 | + |
| 79 | + |
| 80 | +@ModelBase.register("MuseGlimmerForConditionalGeneration") |
| 81 | +class MuseGlimmerVisionModel(MmprojModel): |
| 82 | + def get_vision_config(self) -> dict[str, Any] | None: |
| 83 | + c = self.global_config.get("vision_config") |
| 84 | + if not c: |
| 85 | + return None |
| 86 | + # MuseGlimmer actually uses dynamic size, initialize with nominal size |
| 87 | + image_size = c["pos_emb_height"] * c["patch_size"] * c["merge_size"] |
| 88 | + return {**c, "image_size": image_size} |
| 89 | + |
| 90 | + def set_gguf_parameters(self): |
| 91 | + super().set_gguf_parameters() |
| 92 | + assert self.hparams_vision is not None |
| 93 | + c = self.hparams_vision # enriched vision_config from get_vision_config() |
| 94 | + |
| 95 | + self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.MUSE_GLIMMER) |
| 96 | + self.gguf_writer.add_vision_attention_layernorm_eps(float(c["layer_norm_eps"])) |
| 97 | + self.gguf_writer.add_vision_spatial_merge_size(int(c["merge_size"])) |
| 98 | + |
| 99 | + @classmethod |
| 100 | + def filter_tensors(cls, item): |
| 101 | + name, gen = item |
| 102 | + keep = ("model.vision_tower.", "model.vision_adapter.", "model.vision_projection.") |
| 103 | + if not any(name.startswith(k) for k in keep): |
| 104 | + return None |
| 105 | + return super().filter_tensors((name, gen)) |
| 106 | + |
| 107 | + # 3-layer projector MLP |
| 108 | + _MM_MLP_MAP = { |
| 109 | + "model.vision_adapter.fc1": (gguf.MODEL_TENSOR.V_MMPROJ, 0), |
| 110 | + "model.vision_adapter.fc2": (gguf.MODEL_TENSOR.V_MMPROJ, 1), |
| 111 | + "model.vision_projection": (gguf.MODEL_TENSOR.V_MMPROJ, 2), |
| 112 | + } |
| 113 | + |
| 114 | + def modify_tensors(self, data_torch, name, bid): |
| 115 | + assert self.hparams_vision is not None |
| 116 | + if ".attn.q_proj." in name or ".attn.k_proj." in name: |
| 117 | + n_heads = int(self.hparams_vision["num_attention_heads"]) |
| 118 | + data_torch = _unpermute_for_rope(data_torch, n_heads) |
| 119 | + # Lay out the pt=2 temporal slabs of the patch embedding as a conv2d for build_inp() |
| 120 | + if name.endswith("patch_embedder.patch_embedding.weight"): |
| 121 | + n_embd = data_torch.shape[0] |
| 122 | + pt = int(self.hparams_vision["patch_temporal"]) |
| 123 | + ps = int(self.hparams_vision["patch_size"]) |
| 124 | + data_torch = data_torch.view(n_embd, pt, 3, ps, ps).sum(dim=1) # (n_embd, 3, ps, ps) |
| 125 | + stem, _, suffix = name.rpartition(".") |
| 126 | + if stem in self._MM_MLP_MAP: |
| 127 | + tensor_key, idx = self._MM_MLP_MAP[stem] |
| 128 | + yield (self.format_tensor_name(tensor_key, bid=idx, suffix="." + suffix), data_torch) |
| 129 | + return |
| 130 | + yield (self.map_tensor_name(name), data_torch) |
| 131 | + |
| 132 | + |
| 133 | +@ModelBase.register("MuseGlimmerAssistantModel") |
| 134 | +class MuseGlimmerAssistantModel(TextModel): |
| 135 | + model_arch = gguf.MODEL_ARCH.DFLASH |
| 136 | + |
| 137 | + def set_vocab(self): |
| 138 | + if self.target_model_dir is None: |
| 139 | + raise ValueError( |
| 140 | + "MuseGlimmerAssistant (DFlash drafter) requires --target-model-dir pointing to the " |
| 141 | + "target MuseGlimmer HF directory" |
| 142 | + ) |
| 143 | + |
| 144 | + original_dir = self.dir_model |
| 145 | + self.dir_model = self.target_model_dir |
| 146 | + |
| 147 | + from . import get_model_class |
| 148 | + with open(self.target_model_dir / "config.json", "r", encoding="utf-8") as f: |
| 149 | + target_arch = json.load(f)["architectures"][0] |
| 150 | + target_cls = get_model_class(target_arch) |
| 151 | + if target_cls is not type(self): |
| 152 | + target_cls.set_vocab(self) # ty: ignore[unresolved-attribute] |
| 153 | + else: |
| 154 | + super().set_vocab() |
| 155 | + |
| 156 | + self.dir_model = original_dir |
| 157 | + |
| 158 | + mask_token_id = self.hparams.get("mask_token_id") |
| 159 | + if mask_token_id is not None: |
| 160 | + self.gguf_writer.add_mask_token_id(int(mask_token_id)) |
| 161 | + |
| 162 | + def set_gguf_parameters(self): |
| 163 | + super().set_gguf_parameters() |
| 164 | + h = self.hparams |
| 165 | + |
| 166 | + self.gguf_writer.add_block_size(int(h["block_size"])) |
| 167 | + |
| 168 | + # dflash.target_layers[k] refers to the inputs going into the ith layer, which come from the (i-1)th layer's output. |
| 169 | + # The transformers configuration refers to the outputs being recorded. |
| 170 | + self.gguf_writer.add_target_layers([int(x) + 1 for x in h["target_layer_ids"]]) |
| 171 | + |
| 172 | + if h.get("sliding_window") and h.get("layer_types"): |
| 173 | + self.gguf_writer.add_sliding_window(int(h["sliding_window"])) |
| 174 | + self.gguf_writer.add_sliding_window_pattern([t == "sliding_attention" for t in h["layer_types"]]) |
| 175 | + |
| 176 | + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: |
| 177 | + # DFlash defaults to NEOX (rotate_half) rope, matching transformers HF layout for Q/K, QK-norms |
| 178 | + # no permutation needed. |
| 179 | + yield (self.map_tensor_name(name), data_torch) |
0 commit comments