Felipe97/llama-cpp-compiled

1from __future__ import annotations2 3import re4from collections.abc import Iterable5from typing import TYPE_CHECKING6 7import torch8 9if TYPE_CHECKING:10    from torch import Tensor11 12from .base import ModelBase, TextModel, gguf, logger13 14 15@ModelBase.register("LagunaForCausalLM")16@ModelBase.example("poolside/Laguna-XS.2", "poolside/Laguna-S-2.1")17class LagunaModel(TextModel):18    model_arch = gguf.MODEL_ARCH.LAGUNA19    _experts: list[dict] | None = None20    _gate_types: list[str] | None = None21 22    # --- vocab ---------------------------------------------------------------23 24    def set_vocab(self) -> None:25        self._set_vocab_gpt2()26 27        # Some Laguna releases wrap the chat template in tokenizer_config.json as28        # "{% include 'chat_template.jinja' %}", which SpecialVocab embeds verbatim29        # and llama.cpp's jinja engine cannot process. Prefer the resolved template30        # from the chat_template.jinja file so the GGUF is self-contained.31        tmpl_file = self.dir_model / "chat_template.jinja"32        if tmpl_file.is_file():33            self.gguf_writer.add_chat_template(tmpl_file.read_text(encoding="utf-8"))34            logger.info("gguf: embedded resolved chat_template.jinja (overriding include directive)")35 36        # eos_token_id is a list [2, 24]: token 2 (EOS, also BOS) and token 2437        # (</assistant>, the turn-end). _set_vocab_gpt2 only records the scalar38        # eos, so register the extra id as eot; llama.cpp folds eot into its EOG39        # set, so the model halts on </assistant> natively.40        eos_ids = self.hparams.get("eos_token_id")41        if isinstance(eos_ids, list):42            bos_id = self.hparams.get("bos_token_id")43            extra = [e for e in eos_ids if e != bos_id]44            if extra:45                self.gguf_writer.add_eot_token_id(extra[0])46                logger.info(f"gguf: registered eot_token_id={extra[0]} from eos list {eos_ids}")47 48    def get_vocab_base(self) -> tuple[list[str], list[int], str]:49        # </assistant> is the assistant turn-end (registered as eot below). The50        # HF tokenizer flags it special=false, so the base classifies it as51        # USER_DEFINED and llama.cpp renders its text into generated content,52        # leaking "</assistant>" and breaking response parsing. It is a control53        # marker, so promote it to CONTROL: llama.cpp then treats it as54        # end-of-generation and suppresses its text.55        tokens, toktypes, tokpre = super().get_vocab_base()56        for i, tok in enumerate(tokens):57            if tok == "</assistant>":58                toktypes[i] = gguf.TokenType.CONTROL59                logger.info(f"gguf: marked </assistant> (id {i}) as CONTROL token")60        return tokens, toktypes, tokpre61 62    # --- hparams -------------------------------------------------------------63 64    def set_gguf_parameters(self) -> None:65        super().set_gguf_parameters()66        hparams = self.hparams67 68        # super() does not emit vocab_size for the gpt2 vocab path; head_count is69        # overridden with a per-layer array (XS.2 varies heads per layer via70        # num_attention_heads_per_layer; M.1 is uniform and omits it).71        self.gguf_writer.add_vocab_size(hparams["vocab_size"])72 73        per_layer_heads = hparams.get("num_attention_heads_per_layer")74        if not per_layer_heads:75            per_layer_heads = [hparams["num_attention_heads"]] * hparams["num_hidden_layers"]76        assert len(per_layer_heads) == hparams["num_hidden_layers"], (77            f"num_attention_heads_per_layer length {len(per_layer_heads)} != "78            f"num_hidden_layers {hparams['num_hidden_layers']}"79        )80        self.gguf_writer.add_head_count(per_layer_heads)81 82        # Resolve + validate the attention gate type now so an inconsistent83        # `gating` field fails at conversion time. See _attn_gate_types.84        self._attn_gate_types()85 86        # SWA window size (M.1 has none -> key omitted, swa_type stays NONE).87        sliding_window = hparams.get("sliding_window") or 088        if sliding_window > 0:89            self.gguf_writer.add_sliding_window(sliding_window)90 91        # MoE (expert_count / expert_used_count come from super().set_gguf_parameters())92        self.gguf_writer.add_expert_feed_forward_length(hparams["moe_intermediate_size"])93        self.gguf_writer.add_expert_shared_feed_forward_length(hparams["shared_expert_intermediate_size"])94        self.gguf_writer.add_expert_weights_norm(True)  # HF reference always sum-normalises after top-k95        self.gguf_writer.add_expert_weights_scale(float(hparams["moe_routed_scaling_factor"]))96        self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SIGMOID)97 98        # Leading dense layers (XS.2 has 1, M.1 has 3) before the MoE layers.99        mlp_layer_types: list[str] = hparams["mlp_layer_types"]100        leading_dense = 0101        for t in mlp_layer_types:102            if t == "dense":103                leading_dense += 1104            else:105                break106        self.gguf_writer.add_leading_dense_block_count(leading_dense)107 108        # Per-layer-type RoPE dimension count (partial rotary). base emits109        # rope_freq_base(_swa) and the YaRN params from self.rope_parameters.110        head_dim = hparams["head_dim"]111        full_rope = self.rope_parameters["full_attention"]112        self.gguf_writer.add_rope_dimension_count(113            int(head_dim * float(full_rope.get("partial_rotary_factor", 1.0))))114        swa_rope = self.rope_parameters.get("sliding_attention")115        if swa_rope is not None:116            self.gguf_writer.add_rope_dimension_count_swa(117                int(head_dim * float(swa_rope.get("partial_rotary_factor", 1.0))))118 119    def _attn_gate_types(self) -> list[str]:120        """Per-layer attention output gate type: "per_head" or "per_element".121 122        `gating_types` (per layer) is authoritative when present; otherwise the123        scalar `gating` field is used (the "per-element"/"per-head" string, or124        the legacy boolean True == per-head, as in Laguna-XS.2).125 126        Fails loudly when the model is per-element but the `gating` field does127        not declare that as a string: runtimes that key off `gating` (vLLM,128        transformers) ignore gating_types and read a bare boolean True as129        per-head, silently corrupting the model. Surfacing it here keeps a130        broken checkpoint from being packaged as if it were fine.131        """132        if self._gate_types is not None:133            return self._gate_types134        hparams = self.hparams135        n_layer = hparams["num_hidden_layers"]136        gating = hparams.get("gating")137        gating_types = hparams.get("gating_types")138 139        def _norm(t: object) -> str:140            sval = str(t).replace("-", "_")141            if sval in ("per_element", "per_head"):142                return sval143            raise ValueError(f"Laguna: unrecognised attention gate type {t!r}")144 145        if gating_types:146            assert len(gating_types) == n_layer, (147                f"gating_types length {len(gating_types)} != num_hidden_layers {n_layer}")148            types = [_norm(t) for t in gating_types]149        elif isinstance(gating, str):150            types = [_norm(gating)] * n_layer151        elif gating is True:152            types = ["per_head"] * n_layer153        else:154            raise ValueError(155                f"Laguna: cannot determine attention gate type "156                f"(gating={gating!r}, gating_types={gating_types!r})")157 158        if any(t == "per_element" for t in types) and not (159                isinstance(gating, str) and _norm(gating) == "per_element"):160            raise ValueError(161                f"Laguna config declares a per-element attention gate but "162                f"`gating`={gating!r} is not the string \"per-element\". Runtimes that "163                f"read `gating` (vLLM, transformers) will mis-handle this checkpoint as "164                f"per-head. Set gating=\"per-element\" in the source config.")165 166        self._gate_types = types167        return types168 169    # --- tensor handling -----------------------------------------------------170 171    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:172        # Per-expert MoE weights: model.layers.{bid}.mlp.experts.{xid}.{w}.weight.173        # Only the NUMBERED per-expert weights are stacked; the router bias174        # (mlp.experts.e_score_correction_bias) takes the normal mapping path.175        if re.search(r"mlp\.experts\.\d+\.", name):176            n_experts = self.find_hparam(["num_local_experts", "num_experts"])177            assert bid is not None178            if self._experts is None:179                self._experts = [{} for _ in range(self.block_count)]180            self._experts[bid][name] = data_torch181            needed = [f"model.layers.{bid}.mlp.experts.{x}.{w}.weight"182                      for x in range(n_experts) for w in ("gate_proj", "up_proj", "down_proj")]183            if all(e in self._experts[bid] for e in needed):184                for w_name in ["gate_proj", "up_proj", "down_proj"]:185                    datas = [self._experts[bid][f"model.layers.{bid}.mlp.experts.{x}.{w_name}.weight"]186                             for x in range(n_experts)]187                    stacked = torch.stack(datas, dim=0)188                    merged = f"model.layers.{bid}.mlp.experts.{w_name}.weight"189                    yield from TextModel.modify_tensors(self, stacked, merged, bid)190                self._experts[bid].clear()191                return192            return193        # Cross-check the gate projection width against the declared gate type;194        # a mismatch means the weights and config disagree -> fail, do not guess.195        if bid is not None and name.endswith("self_attn.g_proj.weight"):196            heads = (self.hparams.get("num_attention_heads_per_layer")197                     or [self.hparams["num_attention_heads"]] * self.hparams["num_hidden_layers"])198            n_head = heads[bid]199            head_dim = self.hparams["head_dim"]200            gate_type = self._attn_gate_types()[bid]201            expected = n_head * head_dim if gate_type == "per_element" else n_head202            out_features = int(data_torch.shape[0])203            if out_features != expected:204                raise ValueError(205                    f"Laguna layer {bid}: g_proj output width {out_features} contradicts the "206                    f"declared {gate_type} gate (expected {expected}); weights and config disagree.")207 208        yield from TextModel.modify_tensors(self, data_torch, name, bid)209