char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Sigbjørn Skjæretci : bump ty to 0.0.78 (#28548)67672dc5b

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