char-fork/llama.cpp

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

Xuan-Son Nguyenconvert: add @ModelBase.example (#27208)fa88ae936

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
20.9 KiB466 linesraw
1from __future__ import annotations
2
3import json
4import math
5
6from typing import Callable, Iterable, TYPE_CHECKING
7
8import numpy as np
9import torch
10
11if TYPE_CHECKING:
12    from torch import Tensor
13
14from .base import ModelBase, TextModel, gguf, logger
15
16
17@ModelBase.register(
18    "LLaMAForCausalLM",
19    "LlamaForCausalLM",
20    "MistralForCausalLM",
21    "MixtralForCausalLM",
22    "VLlama3ForCausalLM",
23    "LlavaForConditionalGeneration",
24    "VoxtralForConditionalGeneration",
25    "LlamaForCausalLMEagle3",
26    "Eagle3LlamaForCausalLM",
27    "Eagle3Speculator",
28    "Eagle3DraftModel",
29    "IQuestCoderForCausalLM",
30    "LlamaModel")
31# [TAG_HF_EXAMPLE_GATED] meta-llama/Llama-3.2-1B-Instruct is gated
32@ModelBase.example("unsloth/Llama-3.2-1B-Instruct", "mistralai/Mistral-7B-Instruct-v0.3", "mistralai/Mixtral-8x7B-Instruct-v0.1")
33class LlamaModel(TextModel):
34    model_arch = gguf.MODEL_ARCH.LLAMA
35    undo_permute = True
36
37    def __init__(self, *args, **kwargs):
38        super().__init__(*args, **kwargs)
39        # fix for SmolVLM2, missing `num_attention_heads` in config.json
40        if self.hf_arch == "VLlama3ForCausalLM":
41            self.hparams["num_attention_heads"] = self.hparams.get("num_attention_heads", 32)
42        # Mistral consolidated format has no config.json; origin_hf_arch is HF-only.
43        if self.is_mistral_format:
44            self.origin_hf_arch = None
45        else:
46            hparams = ModelBase.load_hparams(self.dir_model, is_mistral_format=False)
47            self.origin_hf_arch = hparams.get('architectures', [None])[0]
48
49        # Detect eagle3 draft checkpoint by hparams (some models don't use a distinct HF arch name)
50        if "draft_vocab_size" in self.hparams and self.hparams["num_hidden_layers"] == 1:
51            self.is_eagle3 = True
52            self.model_arch = gguf.MODEL_ARCH.EAGLE3
53            logger.info("Detected EAGLE-3 draft model, switching to EAGLE3 architecture")
54            # Re-initialize tensor_map with eagle3 architecture
55            self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
56            # Update gguf_writer architecture
57            self.gguf_writer.arch = gguf.MODEL_ARCH_NAMES[self.model_arch]
58            self.gguf_writer.add_architecture()
59            if self.target_model_dir is None:
60                raise ValueError(
61                    "EAGLE-3 model requires --target-model-dir to be specified. "
62                    "Please provide the path to the target model directory to read config.json"
63                )
64            # Read both eagle3 raw config and target model config
65            with open(self.dir_model / "config.json", 'r', encoding='utf-8') as f:
66                eagle3_raw_config = json.load(f)
67            with open(self.target_model_dir / "config.json", 'r', encoding='utf-8') as f:
68                target_config = json.load(f)
69
70            if "text_config" in target_config:
71                target_config = {**target_config, **target_config["text_config"]}
72            self.target_vocab_size = target_config["vocab_size"]
73
74            # target_layers: use the eagle3 config's explicit aux hidden-state layer ids
75            # if present, else derive from the target layer count.
76            target_num_layers = target_config["num_hidden_layers"]
77            aux_layer_ids = eagle3_raw_config.get("eagle_aux_hidden_state_layer_ids")
78            if aux_layer_ids:
79                target_layers = aux_layer_ids
80            else:
81                target_layers = [2, target_num_layers // 2, target_num_layers - 3]
82            logger.info(f"EAGLE-3: target_layers = {target_layers} (target model has {target_num_layers} layers)")
83            self.gguf_writer.add_target_layers(target_layers)
84
85            # target_hidden_size: prefer eagle3 config, fallback to target config
86            if eagle3_raw_config.get("target_hidden_size") is not None:
87                target_hidden_size = eagle3_raw_config["target_hidden_size"]
88                src = "EAGLE-3 config"
89            else:
90                target_hidden_size = target_config["hidden_size"]
91                src = "target model config"
92            logger.info(f"EAGLE-3: target_hidden_size = {target_hidden_size} (from {src})")
93            self.gguf_writer.add_target_hidden_size(target_hidden_size)
94
95            # norm_before_residual (RedHat-style eagle3 specific)
96            norm_before_residual = eagle3_raw_config.get("norm_before_residual", False)
97            logger.info(f"EAGLE-3: norm_before_residual = {norm_before_residual}")
98            self.gguf_writer.add_norm_before_residual(norm_before_residual)
99
100            # norm_before_fc: RMSNorm applied to the fused target features before the
101            # fc projection (e.g. nvidia/gpt-oss-120b-Eagle3-v3)
102            norm_before_fc = eagle3_raw_config.get("norm_before_fc", False)
103            logger.info(f"EAGLE-3: norm_before_fc = {norm_before_fc}")
104            self.gguf_writer.add_norm_before_fc(norm_before_fc)
105
106    def set_vocab(self):
107        # eagle3: use tokenizer from target model if provided
108        original_dir_model = None
109        if getattr(self, 'is_eagle3', False):
110            assert self.target_model_dir is not None
111            logger.info(f"EAGLE-3: Using tokenizer from target model: {self.target_model_dir}")
112            original_dir_model = self.dir_model
113            self.dir_model = self.target_model_dir
114
115        if self.origin_hf_arch == "GlmasrModel":
116            return self._set_vocab_glmedge()
117
118        if self.is_mistral_format:
119            return self._set_vocab_mistral()
120
121        path_tekken_json = self.dir_model / "tekken.json"
122        path_tokenizer_json = self.dir_model / "tokenizer.json"
123        if path_tekken_json.is_file() and not path_tokenizer_json.is_file():
124            return self._set_vocab_mistral()
125
126        tokenizer_config_file = self.dir_model / 'tokenizer_config.json'
127        if tokenizer_config_file.is_file():
128            with open(tokenizer_config_file, "r", encoding="utf-8") as f:
129                tokenizer_config_json = json.load(f)
130                if (add_prefix_space := tokenizer_config_json.get("add_prefix_space")) is not None:
131                    self.gguf_writer.add_add_space_prefix(add_prefix_space)
132                if tokenizer_config_json.get("tokenizer_class") == "HybridDNATokenizer":
133                    return self._set_vocab_hybriddna()
134
135        try:
136            self._set_vocab_sentencepiece()
137        except FileNotFoundError:
138            try:
139                self._set_vocab_llama_hf()
140            except (FileNotFoundError, TypeError):
141                # Llama 3
142                self._set_vocab_gpt2()
143
144        # Apply to CodeLlama only (and ignore for Llama 3 with a vocab size of 128256)
145        if self.hparams.get("vocab_size", 32000) == 32016:
146            special_vocab = gguf.SpecialVocab(
147                self.dir_model, load_merges=False,
148                special_token_types = ['prefix', 'suffix', 'middle', 'eot']
149            )
150            special_vocab._set_special_token("prefix", 32007)
151            special_vocab._set_special_token("suffix", 32008)
152            special_vocab._set_special_token("middle", 32009)
153            special_vocab._set_special_token("eot",    32010)
154            special_vocab.add_to_gguf(self.gguf_writer)
155
156        # Apply to granite small models only
157        if self.hparams.get("vocab_size", 32000) == 49152:
158            self.gguf_writer.add_add_bos_token(False)
159
160        # eagle3: Restore original dir_model
161        if original_dir_model is not None:
162            self.dir_model = original_dir_model
163
164    def set_gguf_parameters(self):
165        super().set_gguf_parameters()
166        hparams = self.hparams
167
168        if not self.is_mistral_format:
169            self.gguf_writer.add_vocab_size(hparams["vocab_size"])
170
171        if (rope_dim := hparams.get("head_dim")) is None:
172            rope_dim = hparams["hidden_size"] // hparams["num_attention_heads"]
173        self.gguf_writer.add_rope_dimension_count(rope_dim)
174
175    @staticmethod
176    def permute(weights: Tensor, n_head: int, n_head_kv: int | None):
177        if n_head_kv is not None and n_head != n_head_kv:
178            n_head = n_head_kv
179        return (weights.reshape(n_head, 2, weights.shape[0] // n_head // 2, *weights.shape[1:])
180                .swapaxes(1, 2)
181                .reshape(weights.shape))
182
183    def _repack_nvfp4(self, name: str, weight: Tensor, scale: Tensor, scale2: Tensor, input_scale: Tensor):
184        # Mirror the BF16 Q/K RoPE permutation site in modify_tensors; the NVFP4 path bypasses it.
185        if self.undo_permute:
186            n_head = self.find_hparam(["n_heads", "num_attention_heads"], optional=True)
187            n_kv_head = self.find_hparam(["n_kv_heads", "num_key_value_heads"], optional=True)
188            if n_head is not None:
189                if name.endswith("q_proj.weight"):
190                    weight = LlamaModel.permute(weight, n_head, n_head)
191                    scale  = LlamaModel.permute(scale, n_head, n_head)
192                elif name.endswith("k_proj.weight"):
193                    weight = LlamaModel.permute(weight, n_head, n_kv_head)
194                    scale  = LlamaModel.permute(scale, n_head, n_kv_head)
195        super()._repack_nvfp4(name, weight, scale, scale2, input_scale)
196
197    _experts: list[dict[str, Tensor]] | None = None
198
199    @classmethod
200    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
201        name, gen = item
202
203        if "text_model." in name:
204            name = name.replace("text_model.", "") # for SmolVLM
205
206        return super().filter_tensors((name, gen))
207
208    def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Callable[[], Tensor]]:
209        tensors = super().index_tensors(remote_hf_model_id)
210
211        # Handle Eagle3Speculator nested config
212        if "transformer_layer_config" in self.hparams:
213            self.hparams = {**self.hparams, **self.hparams["transformer_layer_config"]}
214
215        # eagle3 detection
216        if "draft_vocab_size" in self.hparams and self.hparams["num_hidden_layers"] == 1:
217            logger.info("EAGLE-3: renaming midlayer.* / layers.0.* to model.layers.0.*")
218            new_tensors = {}
219            for name, gen in tensors.items():
220                if name.startswith("midlayer."):
221                    new_name = "model.layers.0." + name[len("midlayer."):]
222                    new_tensors[new_name] = gen
223                elif name.startswith("layers.0."):  # Eagle3Speculator format
224                    new_name = "model." + name
225                    new_tensors[new_name] = gen
226                else:
227                    new_tensors[name] = gen
228            return new_tensors
229
230        return tensors
231
232    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
233        # eagle3: special tensors that bypass standard llama mapping
234        if getattr(self, 'is_eagle3', False):
235            if name == "fc.weight":
236                yield (name, data_torch)
237                return
238            if name == "input_norm.weight":
239                yield (self.format_tensor_name(gguf.MODEL_TENSOR.ENC_OUTPUT_NORM), data_torch)
240                return
241            if name == "d2t":
242                # store for manual int64 handling in prepare_tensors (avoid F32 conversion)
243                if not hasattr(self, '_eagle3_int_tensors'):
244                    self._eagle3_int_tensors = {}
245                self._eagle3_int_tensors[name] = data_torch
246                return
247            if name == "t2d":
248                # not used at runtime, skip
249                return
250            if name.endswith(".hidden_norm.weight"):
251                yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_NORM_2, bid), data_torch)
252                return
253
254        n_head = self.find_hparam(["n_heads", "num_attention_heads"])
255        n_kv_head = self.find_hparam(["n_kv_heads", "num_key_value_heads"])
256
257        if self.hf_arch == "LlamaModel":
258            name = "model." + name
259
260        if self.undo_permute:
261            if name.endswith(("q_proj.weight", "q_proj.bias")):
262                data_torch = LlamaModel.permute(data_torch, n_head, n_head)
263            if name.endswith(("k_proj.weight", "k_proj.bias")):
264                data_torch = LlamaModel.permute(data_torch, n_head, n_kv_head)
265
266        # process the experts separately
267        if name.find("block_sparse_moe.experts") != -1:
268            n_experts = self.hparams["num_local_experts"]
269
270            assert bid is not None
271
272            if self._experts is None:
273                self._experts = [{} for _ in range(self.block_count)]
274
275            self._experts[bid][name] = data_torch
276
277            if len(self._experts[bid]) >= n_experts * 3:
278                # merge the experts into a single 3d tensor
279                for wid in ["w1", "w2", "w3"]:
280                    datas: list[Tensor] = []
281
282                    for xid in range(n_experts):
283                        ename = f"model.layers.{bid}.block_sparse_moe.experts.{xid}.{wid}.weight"
284                        datas.append(self._experts[bid][ename])
285                        del self._experts[bid][ename]
286
287                    data_torch = torch.stack(datas, dim=0)
288
289                    merged_name = f"layers.{bid}.feed_forward.experts.{wid}.weight"
290
291                    yield from super().modify_tensors(data_torch, merged_name, bid)
292                return
293            else:
294                return
295
296        yield from super().modify_tensors(data_torch, name, bid)
297
298    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
299        if rope_params := self.rope_parameters.get("full_attention", self.rope_parameters):
300            if rope_params.get("rope_type", '').lower() == "llama3":
301                base = rope_params.get("rope_theta", 10000.0)
302                if (dim := self.hparams.get("head_dim")) is None:
303                    dim = self.hparams["hidden_size"] // self.hparams["num_attention_heads"]
304                freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
305
306                factor = rope_params.get("factor", 8.0)
307                low_freq_factor = rope_params.get("low_freq_factor", 1.0)
308                high_freq_factor = rope_params.get("high_freq_factor", 4.0)
309                old_context_len = rope_params.get("original_max_position_embeddings", 8192)
310
311                low_freq_wavelen = old_context_len / low_freq_factor
312                high_freq_wavelen = old_context_len / high_freq_factor
313                # assert low_freq_wavelen != high_freq_wavelen # Errors for Llama4
314
315                rope_factors = []
316                for freq in freqs:
317                    wavelen = 2 * math.pi / freq
318                    if wavelen < high_freq_wavelen:
319                        rope_factors.append(1)
320                    elif wavelen > low_freq_wavelen:
321                        rope_factors.append(factor)
322                    else:
323                        smooth = (old_context_len / wavelen - low_freq_factor) / (high_freq_factor - low_freq_factor)
324                        rope_factors.append(1 / ((1 - smooth) / factor + smooth))
325
326                yield (self.format_tensor_name(gguf.MODEL_TENSOR.ROPE_FREQS), torch.tensor(rope_factors, dtype=torch.float32))
327
328    def prepare_tensors(self):
329        # eagle3: collect d2t original dtype before parent converts tensors to F32
330        eagle3_original_dtypes = {}
331        if getattr(self, 'is_eagle3', False):
332            for name, data_torch in self.get_tensors():
333                if name == "d2t":
334                    eagle3_original_dtypes[name] = data_torch.dtype
335
336        super().prepare_tensors()
337
338        # eagle3: write d2t as absolute target token ids
339        if getattr(self, 'is_eagle3', False) and hasattr(self, '_eagle3_int_tensors'):
340            for name, data_torch in self._eagle3_int_tensors.items():
341                old_dtype = eagle3_original_dtypes.get(name, data_torch.dtype)
342                data = data_torch.to(torch.int64).cpu().numpy()
343                if name == "d2t":
344                    data = data.reshape(-1)
345                    data = data + np.arange(data.size, dtype=np.int64)
346                    if np.any((data < 0) | (data >= self.target_vocab_size)):
347                        raise ValueError(f"EAGLE-3 d2t target ids out of range for target vocab size {self.target_vocab_size}")
348                    if np.unique(data).size != data.size:
349                        raise ValueError("EAGLE-3 d2t contains duplicate target ids")
350                data_qtype = gguf.GGMLQuantizationType.I64
351
352                shape_str = f"{{{', '.join(str(n) for n in reversed(data.shape))}}}"
353                logger.info(f"{name + ',':<30} {old_dtype} --> {data_qtype.name}, shape = {shape_str}")
354                self.gguf_writer.add_tensor(name, data, raw_dtype=data_qtype)
355
356        if self._experts is not None:
357            # flatten `list[dict[str, Tensor]]` into `list[str]`
358            experts = [k for d in self._experts for k in d.keys()]
359            if len(experts) > 0:
360                raise ValueError(f"Unprocessed experts: {experts}")
361
362
363@ModelBase.register("ArceeForCausalLM")
364@ModelBase.example("arcee-ai/AFM-4.5B")
365class ArceeModel(LlamaModel):
366    model_arch = gguf.MODEL_ARCH.ARCEE
367
368    def set_gguf_parameters(self):
369        super().set_gguf_parameters()
370        self._try_set_pooling_type()
371
372
373@ModelBase.register(
374    "Llama4ForConditionalGeneration",
375    "Llama4ForCausalLM",
376)
377# [TAG_HF_EXAMPLE_GATED] meta-llama/Llama-4-Scout-17B-16E-Instruct is gated
378@ModelBase.example("unsloth/Llama-4-Scout-17B-16E-Instruct")
379class Llama4Model(LlamaModel):
380    model_arch = gguf.MODEL_ARCH.LLAMA4
381    undo_permute = False
382
383    def __init__(self, *args, **kwargs):
384        super().__init__(*args, **kwargs)
385        # IMPORTANT: the normal "intermediate_size" is renamed to "intermediate_size_mlp", we need to undo this
386        self.hparams["intermediate_size_moe"] = self.hparams["intermediate_size"]
387        self.hparams["intermediate_size"] = self.hparams["intermediate_size_mlp"]
388
389    def set_vocab(self):
390        self._set_vocab_gpt2()
391
392    def set_gguf_parameters(self):
393        super().set_gguf_parameters()
394        self.gguf_writer.add_interleave_moe_layer_step(self.hparams["interleave_moe_layer_step"])
395        self.gguf_writer.add_expert_feed_forward_length(self.hparams["intermediate_size_moe"])
396        if "layer_types" in self.hparams:
397            if all(lt == "full_attention" for lt in self.hparams["layer_types"]):
398                # all layers are full attention (for MobileLLM), disable swa
399                self.gguf_writer.add_sliding_window(0)
400
401    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None):
402        # split the gate_up into gate and up
403        if "gate_up_proj" in name:
404            name_up = name.replace("gate_up_proj", "up_proj.weight")
405            name_gate = name.replace("gate_up_proj", "gate_proj.weight")
406            dim_half = data_torch.shape[-1] // 2
407            gate_proj_weight, up_proj_weight = data_torch.transpose(-1, -2).split(dim_half, dim=-2)
408            yield from super().modify_tensors(gate_proj_weight, name_gate, bid)
409            yield from super().modify_tensors(up_proj_weight, name_up, bid)
410            return
411
412        if name.endswith("down_proj"):
413            name += ".weight"
414            data_torch = data_torch.transpose(-1, -2)
415
416        yield from super().modify_tensors(data_torch, name, bid)
417
418
419@ModelBase.register("LlamaBidirectionalModel")
420@ModelBase.example("nvidia/llama-embed-nemotron-8b")
421class LlamaEmbedNemotronModel(LlamaModel):
422    model_arch = gguf.MODEL_ARCH.LLAMA_EMBED
423
424
425@ModelBase.register("SmolLM3ForCausalLM")
426@ModelBase.example("HuggingFaceTB/SmolLM3-3B")
427class SmolLM3Model(LlamaModel):
428    model_arch = gguf.MODEL_ARCH.SMOLLM3
429
430
431@ModelBase.register("ApertusForCausalLM")
432@ModelBase.example("swiss-ai/Apertus-8B-Instruct-2509")
433class ApertusModel(LlamaModel):
434    model_arch = gguf.MODEL_ARCH.APERTUS
435    undo_permute = False
436
437    _alpha_n = {}
438    _alpha_p = {}
439    _beta = {}
440    _eps = {}
441
442    def modify_tensors(self, data_torch, name, bid):
443        # Handle xIELU activation parameters
444        n_layers = self.hparams["num_hidden_layers"]
445        if name.endswith(".act_fn.alpha_n"):
446            self._alpha_n[bid] = data_torch.to("cpu").float().item()
447            if (len(self._alpha_n) == n_layers):
448                self.gguf_writer.add_xielu_alpha_n([self._alpha_n[k] for k in sorted(self._alpha_n)])
449            return
450        if name.endswith(".act_fn.alpha_p"):
451            self._alpha_p[bid] = data_torch.to("cpu").float().item()
452            if (len(self._alpha_p) == n_layers):
453                self.gguf_writer.add_xielu_alpha_p([self._alpha_p[k] for k in sorted(self._alpha_p)])
454            return
455        if name.endswith(".act_fn.beta"):
456            self._beta[bid] = data_torch.to("cpu").float().item()
457            if (len(self._beta) == n_layers):
458                self.gguf_writer.add_xielu_beta([self._beta[k] for k in sorted(self._beta)])
459            return
460        if name.endswith(".act_fn.eps"):
461            self._eps[bid] = data_torch.to("cpu").float().item()
462            if (len(self._eps) == n_layers):
463                self.gguf_writer.add_xielu_eps([self._eps[k] for k in sorted(self._eps)])
464            return
465
466        yield from super().modify_tensors(data_torch, name, bid)