Brunobkr/llama.cpp_AlgMor24_github

1from __future__ import annotations2 3import json4import math5 6from typing import Callable, Iterable, TYPE_CHECKING7 8import numpy as np9import torch10 11if TYPE_CHECKING:12    from torch import Tensor13 14from .base import ModelBase, TextModel, gguf, logger15 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")31class LlamaModel(TextModel):32    model_arch = gguf.MODEL_ARCH.LLAMA33    undo_permute = True34 35    def __init__(self, *args, **kwargs):36        super().__init__(*args, **kwargs)37        # fix for SmolVLM2, missing `num_attention_heads` in config.json38        if self.hf_arch == "VLlama3ForCausalLM":39            self.hparams["num_attention_heads"] = self.hparams.get("num_attention_heads", 32)40        # Mistral consolidated format has no config.json; origin_hf_arch is HF-only.41        if self.is_mistral_format:42            self.origin_hf_arch = None43        else:44            hparams = ModelBase.load_hparams(self.dir_model, is_mistral_format=False)45            self.origin_hf_arch = hparams.get('architectures', [None])[0]46 47        # Detect eagle3 draft checkpoint by hparams (some models don't use a distinct HF arch name)48        if "draft_vocab_size" in self.hparams and self.hparams["num_hidden_layers"] == 1:49            self.is_eagle3 = True50            self.model_arch = gguf.MODEL_ARCH.EAGLE351            logger.info("Detected EAGLE-3 draft model, switching to EAGLE3 architecture")52            # Re-initialize tensor_map with eagle3 architecture53            self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)54            # Update gguf_writer architecture55            self.gguf_writer.arch = gguf.MODEL_ARCH_NAMES[self.model_arch]56            self.gguf_writer.add_architecture()57            if self.target_model_dir is None:58                raise ValueError(59                    "EAGLE-3 model requires --target-model-dir to be specified. "60                    "Please provide the path to the target model directory to read config.json"61                )62            # Read both eagle3 raw config and target model config63            with open(self.dir_model / "config.json", 'r', encoding='utf-8') as f:64                eagle3_raw_config = json.load(f)65            with open(self.target_model_dir / "config.json", 'r', encoding='utf-8') as f:66                target_config = json.load(f)67 68            if "text_config" in target_config:69                target_config = {**target_config, **target_config["text_config"]}70            self.target_vocab_size = target_config["vocab_size"]71 72            # target_layers: use the eagle3 config's explicit aux hidden-state layer ids73            # if present, else derive from the target layer count.74            target_num_layers = target_config["num_hidden_layers"]75            aux_layer_ids = eagle3_raw_config.get("eagle_aux_hidden_state_layer_ids")76            if aux_layer_ids:77                target_layers = aux_layer_ids78            else:79                target_layers = [2, target_num_layers // 2, target_num_layers - 3]80            logger.info(f"EAGLE-3: target_layers = {target_layers} (target model has {target_num_layers} layers)")81            self.gguf_writer.add_target_layers(target_layers)82 83            # target_hidden_size: prefer eagle3 config, fallback to target config84            if eagle3_raw_config.get("target_hidden_size") is not None:85                target_hidden_size = eagle3_raw_config["target_hidden_size"]86                src = "EAGLE-3 config"87            else:88                target_hidden_size = target_config["hidden_size"]89                src = "target model config"90            logger.info(f"EAGLE-3: target_hidden_size = {target_hidden_size} (from {src})")91            self.gguf_writer.add_target_hidden_size(target_hidden_size)92 93            # norm_before_residual (RedHat-style eagle3 specific)94            norm_before_residual = eagle3_raw_config.get("norm_before_residual", False)95            logger.info(f"EAGLE-3: norm_before_residual = {norm_before_residual}")96            self.gguf_writer.add_norm_before_residual(norm_before_residual)97 98            # norm_before_fc: RMSNorm applied to the fused target features before the99            # fc projection (e.g. nvidia/gpt-oss-120b-Eagle3-v3)100            norm_before_fc = eagle3_raw_config.get("norm_before_fc", False)101            logger.info(f"EAGLE-3: norm_before_fc = {norm_before_fc}")102            self.gguf_writer.add_norm_before_fc(norm_before_fc)103 104    def set_vocab(self):105        # eagle3: use tokenizer from target model if provided106        original_dir_model = None107        if getattr(self, 'is_eagle3', False):108            assert self.target_model_dir is not None109            logger.info(f"EAGLE-3: Using tokenizer from target model: {self.target_model_dir}")110            original_dir_model = self.dir_model111            self.dir_model = self.target_model_dir112 113        if self.origin_hf_arch == "GlmasrModel":114            return self._set_vocab_glmedge()115 116        if self.is_mistral_format:117            return self._set_vocab_mistral()118 119        path_tekken_json = self.dir_model / "tekken.json"120        path_tokenizer_json = self.dir_model / "tokenizer.json"121        if path_tekken_json.is_file() and not path_tokenizer_json.is_file():122            return self._set_vocab_mistral()123 124        tokenizer_config_file = self.dir_model / 'tokenizer_config.json'125        if tokenizer_config_file.is_file():126            with open(tokenizer_config_file, "r", encoding="utf-8") as f:127                tokenizer_config_json = json.load(f)128                if (add_prefix_space := tokenizer_config_json.get("add_prefix_space")) is not None:129                    self.gguf_writer.add_add_space_prefix(add_prefix_space)130                if tokenizer_config_json.get("tokenizer_class") == "HybridDNATokenizer":131                    return self._set_vocab_hybriddna()132 133        try:134            self._set_vocab_sentencepiece()135        except FileNotFoundError:136            try:137                self._set_vocab_llama_hf()138            except (FileNotFoundError, TypeError):139                # Llama 3140                self._set_vocab_gpt2()141 142        # Apply to CodeLlama only (and ignore for Llama 3 with a vocab size of 128256)143        if self.hparams.get("vocab_size", 32000) == 32016:144            special_vocab = gguf.SpecialVocab(145                self.dir_model, load_merges=False,146                special_token_types = ['prefix', 'suffix', 'middle', 'eot']147            )148            special_vocab._set_special_token("prefix", 32007)149            special_vocab._set_special_token("suffix", 32008)150            special_vocab._set_special_token("middle", 32009)151            special_vocab._set_special_token("eot",    32010)152            special_vocab.add_to_gguf(self.gguf_writer)153 154        # Apply to granite small models only155        if self.hparams.get("vocab_size", 32000) == 49152:156            self.gguf_writer.add_add_bos_token(False)157 158        # eagle3: Restore original dir_model159        if original_dir_model is not None:160            self.dir_model = original_dir_model161 162    def set_gguf_parameters(self):163        super().set_gguf_parameters()164        hparams = self.hparams165 166        if not self.is_mistral_format:167            self.gguf_writer.add_vocab_size(hparams["vocab_size"])168 169        if (rope_dim := hparams.get("head_dim")) is None:170            rope_dim = hparams["hidden_size"] // hparams["num_attention_heads"]171        self.gguf_writer.add_rope_dimension_count(rope_dim)172 173    @staticmethod174    def permute(weights: Tensor, n_head: int, n_head_kv: int | None):175        if n_head_kv is not None and n_head != n_head_kv:176            n_head = n_head_kv177        return (weights.reshape(n_head, 2, weights.shape[0] // n_head // 2, *weights.shape[1:])178                .swapaxes(1, 2)179                .reshape(weights.shape))180 181    def _repack_nvfp4(self, name: str, weight: Tensor, scale: Tensor, scale2: Tensor, input_scale: Tensor):182        # Mirror the BF16 Q/K RoPE permutation site in modify_tensors; the NVFP4 path bypasses it.183        if self.undo_permute:184            n_head = self.find_hparam(["n_heads", "num_attention_heads"], optional=True)185            n_kv_head = self.find_hparam(["n_kv_heads", "num_key_value_heads"], optional=True)186            if n_head is not None:187                if name.endswith("q_proj.weight"):188                    weight = LlamaModel.permute(weight, n_head, n_head)189                    scale  = LlamaModel.permute(scale, n_head, n_head)190                elif name.endswith("k_proj.weight"):191                    weight = LlamaModel.permute(weight, n_head, n_kv_head)192                    scale  = LlamaModel.permute(scale, n_head, n_kv_head)193        super()._repack_nvfp4(name, weight, scale, scale2, input_scale)194 195    _experts: list[dict[str, Tensor]] | None = None196 197    @classmethod198    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:199        name, gen = item200 201        if "text_model." in name:202            name = name.replace("text_model.", "") # for SmolVLM203 204        return super().filter_tensors((name, gen))205 206    def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Callable[[], Tensor]]:207        tensors = super().index_tensors(remote_hf_model_id)208 209        # Handle Eagle3Speculator nested config210        if "transformer_layer_config" in self.hparams:211            self.hparams = {**self.hparams, **self.hparams["transformer_layer_config"]}212 213        # eagle3 detection214        if "draft_vocab_size" in self.hparams and self.hparams["num_hidden_layers"] == 1:215            logger.info("EAGLE-3: renaming midlayer.* / layers.0.* to model.layers.0.*")216            new_tensors = {}217            for name, gen in tensors.items():218                if name.startswith("midlayer."):219                    new_name = "model.layers.0." + name[len("midlayer."):]220                    new_tensors[new_name] = gen221                elif name.startswith("layers.0."):  # Eagle3Speculator format222                    new_name = "model." + name223                    new_tensors[new_name] = gen224                else:225                    new_tensors[name] = gen226            return new_tensors227 228        return tensors229 230    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:231        # eagle3: special tensors that bypass standard llama mapping232        if getattr(self, 'is_eagle3', False):233            if name == "fc.weight":234                yield (name, data_torch)235                return236            if name == "input_norm.weight":237                yield (self.format_tensor_name(gguf.MODEL_TENSOR.ENC_OUTPUT_NORM), data_torch)238                return239            if name == "d2t":240                # store for manual int64 handling in prepare_tensors (avoid F32 conversion)241                if not hasattr(self, '_eagle3_int_tensors'):242                    self._eagle3_int_tensors = {}243                self._eagle3_int_tensors[name] = data_torch244                return245            if name == "t2d":246                # not used at runtime, skip247                return248            if name.endswith(".hidden_norm.weight"):249                yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_NORM_2, bid), data_torch)250                return251 252        n_head = self.find_hparam(["n_heads", "num_attention_heads"])253        n_kv_head = self.find_hparam(["n_kv_heads", "num_key_value_heads"])254 255        if self.hf_arch == "LlamaModel":256            name = "model." + name257 258        if self.undo_permute:259            if name.endswith(("q_proj.weight", "q_proj.bias")):260                data_torch = LlamaModel.permute(data_torch, n_head, n_head)261            if name.endswith(("k_proj.weight", "k_proj.bias")):262                data_torch = LlamaModel.permute(data_torch, n_head, n_kv_head)263 264        # process the experts separately265        if name.find("block_sparse_moe.experts") != -1:266            n_experts = self.hparams["num_local_experts"]267 268            assert bid is not None269 270            if self._experts is None:271                self._experts = [{} for _ in range(self.block_count)]272 273            self._experts[bid][name] = data_torch274 275            if len(self._experts[bid]) >= n_experts * 3:276                # merge the experts into a single 3d tensor277                for wid in ["w1", "w2", "w3"]:278                    datas: list[Tensor] = []279 280                    for xid in range(n_experts):281                        ename = f"model.layers.{bid}.block_sparse_moe.experts.{xid}.{wid}.weight"282                        datas.append(self._experts[bid][ename])283                        del self._experts[bid][ename]284 285                    data_torch = torch.stack(datas, dim=0)286 287                    merged_name = f"layers.{bid}.feed_forward.experts.{wid}.weight"288 289                    yield from super().modify_tensors(data_torch, merged_name, bid)290                return291            else:292                return293 294        yield from super().modify_tensors(data_torch, name, bid)295 296    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:297        if rope_params := self.rope_parameters.get("full_attention", self.rope_parameters):298            if rope_params.get("rope_type", '').lower() == "llama3":299                base = rope_params.get("rope_theta", 10000.0)300                if (dim := self.hparams.get("head_dim")) is None:301                    dim = self.hparams["hidden_size"] // self.hparams["num_attention_heads"]302                freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))303 304                factor = rope_params.get("factor", 8.0)305                low_freq_factor = rope_params.get("low_freq_factor", 1.0)306                high_freq_factor = rope_params.get("high_freq_factor", 4.0)307                old_context_len = rope_params.get("original_max_position_embeddings", 8192)308 309                low_freq_wavelen = old_context_len / low_freq_factor310                high_freq_wavelen = old_context_len / high_freq_factor311                # assert low_freq_wavelen != high_freq_wavelen # Errors for Llama4312 313                rope_factors = []314                for freq in freqs:315                    wavelen = 2 * math.pi / freq316                    if wavelen < high_freq_wavelen:317                        rope_factors.append(1)318                    elif wavelen > low_freq_wavelen:319                        rope_factors.append(factor)320                    else:321                        smooth = (old_context_len / wavelen - low_freq_factor) / (high_freq_factor - low_freq_factor)322                        rope_factors.append(1 / ((1 - smooth) / factor + smooth))323 324                yield (self.format_tensor_name(gguf.MODEL_TENSOR.ROPE_FREQS), torch.tensor(rope_factors, dtype=torch.float32))325 326    def prepare_tensors(self):327        # eagle3: collect d2t original dtype before parent converts tensors to F32328        eagle3_original_dtypes = {}329        if getattr(self, 'is_eagle3', False):330            for name, data_torch in self.get_tensors():331                if name == "d2t":332                    eagle3_original_dtypes[name] = data_torch.dtype333 334        super().prepare_tensors()335 336        # eagle3: write d2t as absolute target token ids337        if getattr(self, 'is_eagle3', False) and hasattr(self, '_eagle3_int_tensors'):338            for name, data_torch in self._eagle3_int_tensors.items():339                old_dtype = eagle3_original_dtypes.get(name, data_torch.dtype)340                data = data_torch.to(torch.int64).cpu().numpy()341                if name == "d2t":342                    data = data.reshape(-1)343                    data = data + np.arange(data.size, dtype=np.int64)344                    if np.any((data < 0) | (data >= self.target_vocab_size)):345                        raise ValueError(f"EAGLE-3 d2t target ids out of range for target vocab size {self.target_vocab_size}")346                    if np.unique(data).size != data.size:347                        raise ValueError("EAGLE-3 d2t contains duplicate target ids")348                data_qtype = gguf.GGMLQuantizationType.I64349 350                shape_str = f"{{{', '.join(str(n) for n in reversed(data.shape))}}}"351                logger.info(f"{name + ',':<30} {old_dtype} --> {data_qtype.name}, shape = {shape_str}")352                self.gguf_writer.add_tensor(name, data, raw_dtype=data_qtype)353 354        if self._experts is not None:355            # flatten `list[dict[str, Tensor]]` into `list[str]`356            experts = [k for d in self._experts for k in d.keys()]357            if len(experts) > 0:358                raise ValueError(f"Unprocessed experts: {experts}")359 360 361@ModelBase.register("ArceeForCausalLM")362class ArceeModel(LlamaModel):363    model_arch = gguf.MODEL_ARCH.ARCEE364 365    def set_gguf_parameters(self):366        super().set_gguf_parameters()367        self._try_set_pooling_type()368 369 370@ModelBase.register(371    "Llama4ForConditionalGeneration",372    "Llama4ForCausalLM",373)374class Llama4Model(LlamaModel):375    model_arch = gguf.MODEL_ARCH.LLAMA4376    undo_permute = False377 378    def __init__(self, *args, **kwargs):379        super().__init__(*args, **kwargs)380        # IMPORTANT: the normal "intermediate_size" is renamed to "intermediate_size_mlp", we need to undo this381        self.hparams["intermediate_size_moe"] = self.hparams["intermediate_size"]382        self.hparams["intermediate_size"] = self.hparams["intermediate_size_mlp"]383 384    def set_vocab(self):385        self._set_vocab_gpt2()386 387    def set_gguf_parameters(self):388        super().set_gguf_parameters()389        self.gguf_writer.add_interleave_moe_layer_step(self.hparams["interleave_moe_layer_step"])390        self.gguf_writer.add_expert_feed_forward_length(self.hparams["intermediate_size_moe"])391        if "layer_types" in self.hparams:392            if all(lt == "full_attention" for lt in self.hparams["layer_types"]):393                # all layers are full attention (for MobileLLM), disable swa394                self.gguf_writer.add_sliding_window(0)395 396    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None):397        # split the gate_up into gate and up398        if "gate_up_proj" in name:399            name_up = name.replace("gate_up_proj", "up_proj.weight")400            name_gate = name.replace("gate_up_proj", "gate_proj.weight")401            dim_half = data_torch.shape[-1] // 2402            gate_proj_weight, up_proj_weight = data_torch.transpose(-1, -2).split(dim_half, dim=-2)403            yield from super().modify_tensors(gate_proj_weight, name_gate, bid)404            yield from super().modify_tensors(up_proj_weight, name_up, bid)405            return406 407        if name.endswith("down_proj"):408            name += ".weight"409            data_torch = data_torch.transpose(-1, -2)410 411        yield from super().modify_tensors(data_torch, name, bid)412 413 414@ModelBase.register("LlamaBidirectionalModel")415class LlamaEmbedNemotronModel(LlamaModel):416    model_arch = gguf.MODEL_ARCH.LLAMA_EMBED417 418 419@ModelBase.register("SmolLM3ForCausalLM")420class SmolLM3Model(LlamaModel):421    model_arch = gguf.MODEL_ARCH.SMOLLM3422 423 424@ModelBase.register("ApertusForCausalLM")425class ApertusModel(LlamaModel):426    model_arch = gguf.MODEL_ARCH.APERTUS427    undo_permute = False428 429    _alpha_n = {}430    _alpha_p = {}431    _beta = {}432    _eps = {}433 434    def modify_tensors(self, data_torch, name, bid):435        # Handle xIELU activation parameters436        n_layers = self.hparams["num_hidden_layers"]437        if name.endswith(".act_fn.alpha_n"):438            self._alpha_n[bid] = data_torch.to("cpu").float().item()439            if (len(self._alpha_n) == n_layers):440                self.gguf_writer.add_xielu_alpha_n([self._alpha_n[k] for k in sorted(self._alpha_n)])441            return442        if name.endswith(".act_fn.alpha_p"):443            self._alpha_p[bid] = data_torch.to("cpu").float().item()444            if (len(self._alpha_p) == n_layers):445                self.gguf_writer.add_xielu_alpha_p([self._alpha_p[k] for k in sorted(self._alpha_p)])446            return447        if name.endswith(".act_fn.beta"):448            self._beta[bid] = data_torch.to("cpu").float().item()449            if (len(self._beta) == n_layers):450                self.gguf_writer.add_xielu_beta([self._beta[k] for k in sorted(self._beta)])451            return452        if name.endswith(".act_fn.eps"):453            self._eps[bid] = data_torch.to("cpu").float().item()454            if (len(self._eps) == n_layers):455                self.gguf_writer.add_xielu_eps([self._eps[k] for k in sorted(self._eps)])456            return457 458        yield from super().modify_tensors(data_torch, name, bid)459