char-fork/llama.cpp

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

jacekpoplawskiconvert: fix GLM regression in index_tensors (#27655)71cc86fa4

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
18.0 KiB404 linesraw
1from __future__ import annotations
2
3import re
4
5from typing import Callable, Iterable, TYPE_CHECKING
6
7import torch
8
9if TYPE_CHECKING:
10    from torch import Tensor
11
12from .base import ModelBase, TextModel, gguf, logger
13
14from .deepseek import DeepseekV2Model
15
16
17@ModelBase.register("Glm4ForCausalLM", "Glm4vForConditionalGeneration")
18@ModelBase.example("zai-org/GLM-4-9B-0414")
19class Glm4Model(TextModel):
20    model_arch = gguf.MODEL_ARCH.GLM4
21    use_mrope = False
22    partial_rotary_factor = 0.5
23
24    def __init__(self, *args, **kwargs):
25        super().__init__(*args, **kwargs)
26        self.partial_rotary_factor = self.rope_parameters.get("partial_rotary_factor", 0.5)
27        if "mrope_section" in self.rope_parameters:
28            self.use_mrope = True
29            logger.info("Q/K weight will need to be permuted for M-RoPE")
30
31    def set_vocab(self):
32        from transformers import AutoTokenizer
33        tokenizer = AutoTokenizer.from_pretrained(self.dir_model, trust_remote_code=True)
34        special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True)
35        tokens, toktypes, tokpre = self.get_vocab_base()
36        self.gguf_writer.add_tokenizer_model("gpt2")
37        self.gguf_writer.add_tokenizer_pre(tokpre)
38        self.gguf_writer.add_token_list(tokens)
39        self.gguf_writer.add_token_types(toktypes)
40        special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True)
41        special_vocab._set_special_token("eos", tokenizer.get_added_vocab()["<|endoftext|>"])  # ty: ignore[unresolved-attribute]
42        special_vocab._set_special_token("eot", tokenizer.get_added_vocab()["<|user|>"])  # ty: ignore[unresolved-attribute]
43        special_vocab._set_special_token("unk", tokenizer.get_added_vocab()["<|endoftext|>"])  # ty: ignore[unresolved-attribute]
44        special_vocab._set_special_token("bos", tokenizer.get_added_vocab()["<|endoftext|>"])  # ty: ignore[unresolved-attribute]
45        special_vocab.add_to_gguf(self.gguf_writer)
46
47    def set_gguf_parameters(self):
48        super().set_gguf_parameters()
49        if (rope_dim := self.hparams.get("head_dim")) is None:
50            rope_dim = self.hparams["hidden_size"] // self.hparams["num_attention_heads"]
51        self.gguf_writer.add_rope_dimension_count(int(rope_dim * self.partial_rotary_factor))
52
53    @staticmethod
54    def normal_to_neox(weights: Tensor, n_head: int, n_head_kv: int, head_dim: int, partial_rotary_factor: float) -> Tensor:
55        orig_shape = weights.shape
56        if len(orig_shape) == 1:
57            weights = weights.unsqueeze(1)  # [out_dim, 1]
58        if len(weights.shape) != 2:
59            raise ValueError("Only 1D and 2D tensors are supported.")
60        n_effective_heads = weights.shape[0] // head_dim
61        if n_head_kv is not None and n_effective_heads != n_head:
62            if n_effective_heads != n_head_kv:
63                raise AssertionError(f"Mismatch in effective heads: computed {n_effective_heads}, expected {n_head} or {n_head_kv}")
64        rotary_dim = int(head_dim * partial_rotary_factor)
65        if rotary_dim % 2 != 0:
66            raise ValueError("rotary_dim must be even.")
67        reshaped = weights.reshape(n_effective_heads, head_dim, -1)
68        rot_part = reshaped[:, :rotary_dim, :]
69        non_rot_part = reshaped[:, rotary_dim:, :]
70        permuted_rot = torch.cat((rot_part[:, ::2, :], rot_part[:, 1::2, :]), dim=1)
71        combined = torch.cat((permuted_rot, non_rot_part), dim=1)
72        result = combined.reshape(weights.shape)
73        return result if len(orig_shape) != 1 else result.squeeze(1)
74
75    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
76        if self.use_mrope:
77            n_head = self.hparams["num_attention_heads"]
78            n_kv_head = self.hparams["num_key_value_heads"]
79            n_embd = self.hparams["hidden_size"]
80            head_dim = self.hparams.get("head_dim", n_embd // n_head)
81            # because llama.cpp M-RoPE kernel only supports Neox ordering, we have to permute the weights here
82            if name.endswith(("q_proj.weight", "q_proj.bias")):
83                data_torch = Glm4Model.normal_to_neox(data_torch, n_head, n_head, head_dim, self.partial_rotary_factor)
84            if name.endswith(("k_proj.weight", "k_proj.bias")):
85                data_torch = Glm4Model.normal_to_neox(data_torch, n_head, n_kv_head, head_dim, self.partial_rotary_factor)
86        yield from super().modify_tensors(data_torch, name, bid)
87
88
89@ModelBase.register("GlmOcrForConditionalGeneration")
90@ModelBase.example("zai-org/GLM-OCR")
91class GlmOCRModel(Glm4Model):
92    model_arch = gguf.MODEL_ARCH.GLM4
93    use_mrope = False
94    partial_rotary_factor = 0.5
95
96    # Note: GLM-OCR is the same as GLM4, but with an extra NextN/MTP prediction layer
97
98    def __init__(self, *args, **kwargs):
99        super().__init__(*args, **kwargs)
100        # GLM-OCR has num_hidden_layers + 1 actual layers (including NextN layer)
101        self.block_count = self.hparams["num_hidden_layers"] + self.hparams.get("num_nextn_predict_layers", 0)
102        self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
103
104    def set_gguf_parameters(self):
105        super().set_gguf_parameters()
106        # NextN/MTP prediction layers
107        if (num_nextn_predict_layers := self.hparams.get("num_nextn_predict_layers")) is not None:
108            self.gguf_writer.add_nextn_predict_layers(num_nextn_predict_layers)
109
110
111@ModelBase.register("Glm4MoeForCausalLM", "Glm4vMoeForConditionalGeneration")
112@ModelBase.example("zai-org/GLM-4.5-Air")
113class Glm4MoeModel(TextModel):
114    model_arch = gguf.MODEL_ARCH.GLM4_MOE
115    supports_mtp_export = True
116    _n_main_layers: int | None = None
117
118    def __init__(self, *args, **kwargs):
119        super().__init__(*args, **kwargs)
120        if not self.no_mtp:
121            self.block_count += self.hparams.get("num_nextn_predict_layers", 0)
122            self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
123
124    def index_tensors(self, remote_hf_model_id: str | None = None):
125        hparams = {**self.hparams, **self.hparams.get("text_config", {})}
126        key = next((k for k in ["n_layers", "num_hidden_layers", "n_layer", "num_layers"] if k in hparams), None)
127        type(self)._n_main_layers = hparams.get(key)
128        return super().index_tensors(remote_hf_model_id=remote_hf_model_id)
129
130    @classmethod
131    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
132        if (titem := super().filter_tensors(item)) is None:
133            return None
134        name, gen = titem
135
136        assert cls._n_main_layers is not None
137        is_mtp = (m := re.match(r"model\.layers\.(\d+)\.", name)) is not None and int(m.group(1)) >= cls._n_main_layers
138
139        if is_mtp and cls.no_mtp:
140            return None
141        if cls.mtp_only and not is_mtp and name not in (
142            "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight",
143        ):
144            return None
145
146        return name, gen
147
148    def set_vocab(self):
149        return self._set_vocab_glm()
150
151    def set_gguf_parameters(self):
152        super().set_gguf_parameters()
153        if (rope_dim := self.hparams.get("head_dim")) is None:
154            rope_dim = (
155                self.hparams["hidden_size"] // self.hparams["num_attention_heads"]
156            )
157        self.gguf_writer.add_rope_dimension_count(
158            int(rope_dim * self.rope_parameters.get("partial_rotary_factor", 0.5))
159        )
160
161        # MoE parameters - Use only routed expert count (shared experts handled separately)
162        if (n_routed_experts := self.hparams.get("n_routed_experts")) is not None:
163            self.gguf_writer.add_expert_count(n_routed_experts)
164        if (moe_intermediate_size := self.hparams.get("moe_intermediate_size")) is not None:
165            self.gguf_writer.add_expert_feed_forward_length(moe_intermediate_size)
166        if (n_shared_experts := self.hparams.get("n_shared_experts")) is not None:
167            self.gguf_writer.add_expert_shared_count(n_shared_experts)
168        if (first_k_dense_replace := self.hparams.get("first_k_dense_replace")) is not None:
169            self.gguf_writer.add_leading_dense_block_count(first_k_dense_replace)
170
171        # Expert gating function (sigmoid for GLM4_MOE)
172        self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SIGMOID)
173
174        # Routed scaling factor
175        if (routed_scaling_factor := self.hparams.get("routed_scaling_factor")) is not None:
176            self.gguf_writer.add_expert_weights_scale(routed_scaling_factor)
177
178        # Normalise topk probabilities
179        if (norm_topk_prob := self.hparams.get("norm_topk_prob")) is not None:
180            self.gguf_writer.add_expert_weights_norm(norm_topk_prob)
181
182        if not self.no_mtp and (num_nextn_predict_layers := self.hparams.get("num_nextn_predict_layers")) is not None:
183            self.gguf_writer.add_nextn_predict_layers(num_nextn_predict_layers)
184
185    def prepare_metadata(self, vocab_only: bool):
186        from_dir = self.fname_out.is_dir()
187        super().prepare_metadata(vocab_only=vocab_only)
188
189        if not self.mtp_only or not from_dir:
190            return
191
192        output_type: str = self.ftype.name.partition("_")[2]
193        fname_default: str = gguf.naming_convention(
194            self.metadata.name, self.metadata.basename, self.metadata.finetune,
195            self.metadata.version, size_label=None, output_type=output_type, model_type=None)
196        self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf"
197
198    _experts: list[dict[str, Tensor]] | None = None
199
200    # note: unlike GLM4V non-MoE, we don't need to permute Q/K here since GLM4V_MOE uses Neox ordering already
201    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
202        # Handle main token embedding (but not layer-specific NextN embeddings)
203        if name == "model.embed_tokens.weight" and ".layers." not in name:
204            yield from super().modify_tensors(data_torch, "token_embd.weight", bid)
205            return
206
207        # Handle routed experts
208        if name.find("mlp.experts") != -1:
209            n_experts = self.hparams["n_routed_experts"]
210            assert bid is not None
211
212            if self._experts is None:
213                self._experts = [{} for _ in range(self.block_count)]
214
215            self._experts[bid][name] = data_torch
216
217            if len(self._experts[bid]) >= n_experts * 3:
218                # merge the experts into a single 3d tensor
219                for w_name in ["down_proj", "gate_proj", "up_proj"]:
220                    datas: list[Tensor] = []
221
222                    for xid in range(n_experts):
223                        ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"
224                        datas.append(self._experts[bid][ename])
225                        del self._experts[bid][ename]
226
227                    data_torch = torch.stack(datas, dim=0)
228
229                    merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"
230
231                    yield from super().modify_tensors(data_torch, merged_name, bid)
232                return
233            else:
234                return
235
236        yield from super().modify_tensors(data_torch, name, bid)
237
238    def prepare_tensors(self):
239        super().prepare_tensors()
240        if self._experts is not None:
241            # flatten `list[dict[str, Tensor]]` into `list[str]`
242            experts = [k for d in self._experts for k in d.keys()]
243            if len(experts) > 0:
244                raise ValueError(f"Unprocessed experts: {experts}")
245
246
247@ModelBase.register("Glm4MoeLiteForCausalLM")
248@ModelBase.example("zai-org/GLM-4.7-Flash")
249class Glm4MoeLiteModel(DeepseekV2Model):
250    model_arch = gguf.MODEL_ARCH.DEEPSEEK2
251    skip_mtp = False
252    supports_mtp_export = True
253    _n_main_layers: int | None = None
254
255    def set_vocab(self):
256        return self._set_vocab_glm()
257
258    def __init__(self, *args, **kwargs):
259        super().__init__(*args, **kwargs)
260
261        num_hidden_layers = self.hparams["num_hidden_layers"]
262        self.num_nextn_predict_layers = self.hparams.get("num_nextn_predict_layers", 0)
263        self.skip_mtp = self.no_mtp or self.num_nextn_predict_layers == 0
264
265        if self.skip_mtp:
266            self.block_count = num_hidden_layers
267        else:
268            self.block_count = num_hidden_layers + self.num_nextn_predict_layers
269
270        self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
271
272    def set_gguf_parameters(self):
273        super().set_gguf_parameters()
274
275        if self.skip_mtp:
276            return
277
278        self.gguf_writer.add_nextn_predict_layers(self.num_nextn_predict_layers)
279
280    def index_tensors(self, remote_hf_model_id: str | None = None):
281        type(self)._n_main_layers = self.hparams["num_hidden_layers"]
282        return super().index_tensors(remote_hf_model_id=remote_hf_model_id)
283
284    @classmethod
285    def filter_tensors(cls, item):
286        if (titem := super().filter_tensors(item)) is None:
287            return None
288        name, gen = titem
289
290        if cls._n_main_layers is not None:
291            match = re.match(r"model\.layers\.(\d+)\.", name)
292            is_mtp = match is not None and int(match.group(1)) >= cls._n_main_layers
293            if is_mtp and cls.no_mtp:
294                return None
295            if cls.mtp_only and not is_mtp and name not in (
296                "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight",
297            ):
298                return None
299
300        return name, gen
301
302    def prepare_metadata(self, vocab_only: bool):
303        from_dir = self.fname_out.is_dir()
304        super().prepare_metadata(vocab_only=vocab_only)
305
306        if not self.mtp_only or not from_dir:
307            return
308
309        output_type: str = self.ftype.name.partition("_")[2]
310        fname_default: str = gguf.naming_convention(
311            self.metadata.name, self.metadata.basename, self.metadata.finetune,
312            self.metadata.version, size_label=None, output_type=output_type, model_type=None)
313        self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf"
314
315
316@ModelBase.register("GlmMoeDsaForCausalLM")
317@ModelBase.example("zai-org/GLM-5.2")
318class GlmMoeDsaModel(DeepseekV2Model):
319    model_arch = gguf.MODEL_ARCH.GLM_DSA
320    skip_mtp = False
321    supports_mtp_export = True
322
323    # Trunk layer count, stashed before indexing so the classmethod
324    # filter_tensors can identify the appended NextN/MTP block (mirrors
325    # HYV3Model / Step35Model).
326    _n_main_layers: int | None = None
327
328    def __init__(self, *args, **kwargs):
329        super().__init__(*args, **kwargs)
330        self.block_count = self.hparams["num_hidden_layers"]
331        if not self.no_mtp:
332            self.block_count += self.hparams.get("num_nextn_predict_layers", 0)
333        self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
334
335    def index_tensors(self, remote_hf_model_id: str | None = None):
336        type(self)._n_main_layers = self.hparams["num_hidden_layers"]
337        return super().index_tensors(remote_hf_model_id=remote_hf_model_id)
338
339    @classmethod
340    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
341        if (titem := super().filter_tensors(item)) is None:
342            return None
343        name, gen = titem
344
345        # GLM-5.2 appends the NextN/MTP block past num_hidden_layers
346        # (model.layers.78 -> blk.78 in the 79-block file).
347        assert cls._n_main_layers is not None
348        is_mtp = (m := re.match(r"model\.layers\.(\d+)\.", name)) is not None and int(m.group(1)) >= cls._n_main_layers
349
350        # --no-mtp: drop the appended NextN block entirely.
351        if is_mtp and cls.no_mtp:
352            return None
353        # --mtp: keep ONLY NextN-block tensors plus the shared embeddings/
354        # norm/lm_head (so the resulting GGUF carries just the draft head).
355        if cls.mtp_only and not is_mtp and name not in (
356            "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight",
357        ):
358            return None
359
360        return name, gen
361
362    def set_vocab(self):
363        return self._set_vocab_glm()
364
365    def set_gguf_parameters(self):
366        super().set_gguf_parameters()
367
368        rope_dim = self.hparams["qk_rope_head_dim"]
369        partial_rotary_factor = self.rope_parameters.get("partial_rotary_factor", 1.0)
370        self.gguf_writer.add_rope_dimension_count(int(rope_dim * partial_rotary_factor))
371
372        # NextN/MTP prediction layers
373        if not self.no_mtp and (num_nextn_predict_layers := self.hparams.get("num_nextn_predict_layers")) is not None:
374            self.gguf_writer.add_nextn_predict_layers(num_nextn_predict_layers)
375
376        # DSA indexer parameters
377        self.gguf_writer.add_indexer_head_count(self.hparams["index_n_heads"])
378        self.gguf_writer.add_indexer_key_length(self.hparams["index_head_dim"])
379        self.gguf_writer.add_indexer_top_k(self.hparams["index_topk"])
380        if (indexer_types := self.hparams.get("indexer_types")) is not None:
381            indexer_types = [t == "full" for t in indexer_types]
382            self.gguf_writer.add_indexer_types(indexer_types)
383
384
385@ModelBase.register("SolarOpenForCausalLM")
386@ModelBase.example("upstage/Solar-Open-100B")
387class SolarOpenModel(Glm4MoeModel):
388    model_arch = gguf.MODEL_ARCH.GLM4_MOE
389    supports_mtp_export = False
390
391    def set_vocab(self):
392        from transformers import AutoTokenizer
393        tokenizer = AutoTokenizer.from_pretrained(self.dir_model)
394        special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True)
395        tokens, toktypes, tokpre = self.get_vocab_base()
396        self.gguf_writer.add_tokenizer_model("gpt2")
397        self.gguf_writer.add_tokenizer_pre(tokpre)
398        self.gguf_writer.add_token_list(tokens)
399        self.gguf_writer.add_token_types(toktypes)
400        special_vocab._set_special_token("eos", tokenizer.get_added_vocab()["<|endoftext|>"])  # ty: ignore[unresolved-attribute]
401        special_vocab._set_special_token("eot", tokenizer.get_added_vocab()["<|endoftext|>"])  # ty: ignore[unresolved-attribute]
402        special_vocab._set_special_token("unk", tokenizer.get_added_vocab()["<unk>"])  # ty: ignore[unresolved-attribute]
403        special_vocab._set_special_token("bos", tokenizer.get_added_vocab()["<|startoftext|>"])  # ty: ignore[unresolved-attribute]
404        special_vocab.add_to_gguf(self.gguf_writer)