Buckets:
| 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 | |
| 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) | |
Xet Storage Details
- Size:
- 8.67 kB
- Xet hash:
- 5703f07b7cc242f1a751313a2a532145aa4a9ef1ad7ddf5268bdf5b0f2e89b6a
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.