from __future__ import annotations from typing import Iterable, TYPE_CHECKING import torch if TYPE_CHECKING: from torch import Tensor from .base import ModelBase, TextModel, gguf @ModelBase.register("LimiteForCausalLM") @ModelBase.example("paradigma-inc/limite-1b-violetto") class LimiteModel(TextModel): model_arch = gguf.MODEL_ARCH.LIMITE def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) hp = self.hparams expected = { "pos_mode": "rope", "qk_norm": "rms_pre_rope", "rope_style": "interleaved_pairs_odd_lane_sign_flip", "rope_per_layer": False, "rms_norm_has_weight": False, "rms_norm_eps_mode": "torch_finfo_default", "mlp_type": "swiglu", "attn_gate_applied": "per_head_before_o_proj", "ve_head_slice": "first_num_key_value_heads", "ve_applied_before_qk_norm": True, "global_nope": True, "global_window": -1, "mudd_r_site": "resid", "tie_word_embeddings": True, "final_softcap": 0.0, } for key, value in expected.items(): if hp.get(key) != value: raise NotImplementedError(f"Limite: unsupported {key}={hp.get(key)!r}, expected {value!r}") if hp["softcap_logits"]["kind"] != "sigmoid": raise NotImplementedError(f"Limite: unsupported softcap_logits kind {hp['softcap_logits']['kind']!r}") self.n_head = hp["num_attention_heads"] self.n_head_kv = hp["num_key_value_heads"] self.head_dim = hp["head_dim"] self.n_pairs = hp["rope_n_pairs"] self.mudd_taps: dict[int, list[int]] = {} if hp.get("mudd"): if not hp.get("mudd_mlp"): raise NotImplementedError("Limite: mudd without mudd_mlp is not supported") self.mudd_taps = {int(il): [int(t) for t in taps] for il, taps in hp["mudd_tap_idx"].items()} assert all(len(taps) == hp["mudd_taps"] for taps in self.mudd_taps.values()) def set_vocab(self): self._set_vocab_gpt2() def set_gguf_parameters(self): super().set_gguf_parameters() hp = self.hparams self.gguf_writer.add_key_length(self.head_dim) self.gguf_writer.add_value_length(self.head_dim) # torch rms_norm with eps=None uses the fp32 epsilon, even for bf16 inputs self.gguf_writer.add_layer_norm_rms_eps(torch.finfo(torch.float32).eps) self.gguf_writer.add_attention_scale(hp["attention_softmax_scale"]) self.gguf_writer.add_xsa_eps(hp["xsa_normalize_eps"] if hp.get("xsa") else 0.0) # HF freq_i = base^(-i/(n_pairs-1)), ggml uses freq_base^(-i/n_pairs) self.gguf_writer.add_rope_dimension_count(2 * self.n_pairs) self.gguf_writer.add_rope_freq_base(hp["rope_base_local"] ** (self.n_pairs / (self.n_pairs - 1))) # HF keeps keys k >= q - window, so a window of W spans W + 1 tokens global_layers = set(hp["global_layers"]) self.gguf_writer.add_sliding_window(hp["sliding_window"] + 1) self.gguf_writer.add_sliding_window_pattern([il not in global_layers for il in range(self.block_count)]) cap = hp["softcap_logits"] self.gguf_writer.add_final_logit_sigmoid_capping(cap["a"], cap["b"], cap["c"]) if self.mudd_taps: n_taps = hp["mudd_taps"] tap_indices = [-1] * (self.block_count * n_taps) for il, taps in self.mudd_taps.items(): tap_indices[il * n_taps:(il + 1) * n_taps] = taps self.gguf_writer.add_mudd_feed_forward_length(hp["mudd_inter"]) self.gguf_writer.add_mudd_tap_count(n_taps) self.gguf_writer.add_mudd_tap_indices(tap_indices) def tensor_force_quant(self, name, new_name, bid, n_dims): # tiny mixing weights, kept at full precision as in the reference # the mudd tensors come from a global source tensor, so take the block id from the new name if new_name.startswith("blk."): bid = int(new_name.split(".")[1]) for key in (gguf.MODEL_TENSOR.ATTN_GATE, gguf.MODEL_TENSOR.ATTN_VE_GATE, gguf.MODEL_TENSOR.MUDD_DOWN, gguf.MODEL_TENSOR.MUDD_UP): if self.match_model_tensor_name(new_name, key, bid): return gguf.GGMLQuantizationType.F32 return super().tensor_force_quant(name, new_name, bid, n_dims) def _swap_rope_pairs(self, w: Tensor, n_head: int) -> Tensor: # HF rotates each pair by -theta; swapping the pair lanes gives the ggml +theta rotation. # q and k are both swapped, so the attention scores do not change. idx = torch.arange(self.head_dim) idx[:2 * self.n_pairs] = idx[:2 * self.n_pairs].view(-1, 2).flip(-1).reshape(-1) w = w.reshape(n_head, self.head_dim, -1) return w[:, idx, :].reshape(n_head * self.head_dim, -1) def _pad_cols(self, w: Tensor) -> Tensor: # gates read only the first channels of the input, pad them to n_embd return torch.nn.functional.pad(w, (0, self.hparams["hidden_size"] - w.shape[-1])) def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: hp = self.hparams T = gguf.MODEL_TENSOR def load(n: str) -> Tensor: return self.model_tensors[n]().float() if name == "model.value_embeds.weight": ve = data_torch.reshape(data_torch.shape[0], -1, hp["ve_dim"])[:, :self.n_head_kv] ve = torch.nn.functional.pad(ve, (0, self.head_dim - hp["ve_dim"])) yield self.format_tensor_name(T.VALUE_EMBD), ve.reshape(ve.shape[0], -1) * hp["ve_gate_scale"] return if name == "model.mudd.dense1": for il in self.mudd_taps: yield self.format_tensor_name(T.MUDD_DOWN, il), data_torch return if name == "model.mudd.dense2": # rows [0, n) mix the attention input, rows [n, 2n) mix the residual dense2_mlp = load("model.mudd.dense2_mlp") bias = load("model.mudd.bias") bias_mlp = load("model.mudd.bias_mlp") for il, taps in self.mudd_taps.items(): n = len(taps) w = torch.cat([data_torch[il, :n], dense2_mlp[il, :n]], 0) b = torch.cat([bias[il, :n], bias_mlp[il, :n]], 0) yield self.format_tensor_name(T.MUDD_UP, il), w yield self.format_tensor_name(T.MUDD_UP, il, ".bias"), b return if name.startswith("model.mudd."): return if bid is None: yield from super().modify_tensors(data_torch, name, bid) return prefix = f"model.layers.{bid}." name = name.removeprefix(prefix) if name in ("self_attn.q_proj.weight", "self_attn.k_proj.weight", "self_attn.v_proj.weight"): data_torch = data_torch * load(prefix + "self_attn.qkv_scale") if name == "self_attn.q_proj.weight": data_torch = self._swap_rope_pairs(data_torch, self.n_head) elif name == "self_attn.k_proj.weight": data_torch = self._swap_rope_pairs(data_torch, self.n_head_kv) elif name == "self_attn.o_proj.weight": data_torch = data_torch * load(prefix + "self_attn.o_scale") * load(prefix + "post_lambda_attn") * hp["attn_gate_scale"] elif name == "mlp.down_proj.weight": data_torch = data_torch * load(prefix + "post_lambda_mlp") elif name == "self_attn.attn_gate": yield self.format_tensor_name(T.ATTN_GATE, bid), self._pad_cols(data_torch) return elif name == "self_attn.ve_gate": yield self.format_tensor_name(T.ATTN_VE_GATE, bid), self._pad_cols(data_torch[:self.n_head_kv]) return elif name == "self_attn.xsa_alpha": if hp.get("xsa") and bid in hp["xsa_layers"]: yield self.format_tensor_name(T.ATTN_XSA_ALPHA, bid), torch.tanh(data_torch) return elif name == "resid_lambda_attn": yield self.format_tensor_name(T.ATTN_RESID_SCALE, bid), data_torch.reshape(1) return elif name == "resid_lambda_mlp": yield self.format_tensor_name(T.FFN_RESID_SCALE, bid), data_torch.reshape(1) return elif name in ("self_attn.qkv_scale", "self_attn.o_scale", "post_lambda_attn", "post_lambda_mlp"): # folded into the projections above return yield from super().modify_tensors(data_torch, prefix + name, bid)