configuration_limite.py

"""Configuration and validation for Limite checkpoints.

Every numerical choice is explicit in ``config.json``. Unsupported values fail
at load time instead of silently selecting different model behavior.

Field semantics are documented in the artifact's own `config_field_notes.json`.
"""

from typing import Any

from transformers.configuration_utils import PretrainedConfig

from .contract import ARCHITECTURE, BOS_TOKEN_ID, EOS_TOKEN_ID, MODEL_TYPE, PAD_TOKEN_ID

REQUIRED_CONFIG_FIELDS: tuple[str, ...] = (
    "architectures",
    "model_type",
    "hidden_size",
    "num_hidden_layers",
    "num_attention_heads",
    "num_key_value_heads",
    "head_dim",
    "intermediate_size",
    "vocab_size",
    "padded_vocab_size",
    "tokenizer_vocab_size",
    "max_position_embeddings",
    "tie_word_embeddings",
    "torch_dtype",
    "rms_norm_has_weight",
    "rms_norm_eps_mode",
    "qk_norm",
    "attention_softmax_scale",
    "sliding_window",
    "sliding_window_convention",
    "global_window",
    "global_layers",
    "global_every",
    "global_nope",
    "attn_gate_channels",
    "attn_gate_scale",
    "attn_gate_applied",
    "pos_mode",
    "rope_frac",
    "rope_base_local",
    "rope_base_global",
    "rope_per_layer",
    "rope_n_pairs",
    "rope_style",
    "rope_cos_sin_dtype",
    "ve_dim",
    "ve_layers",
    "ve_gate_channels",
    "ve_gate_scale",
    "ve_head_slice",
    "ve_stored_heads",
    "ve_applied_before_qk_norm",
    "xsa",
    "xsa_layers",
    "xsa_normalize_eps",
    "mudd",
    "mudd_at",
    "mudd_layers",
    "mudd_taps",
    "mudd_inter",
    "mudd_tap_idx",
    "mudd_mlp",
    "mudd_hist_convention",
    "mudd_accumulation",
    "mlp_type",
    "mlp_formula",
    "mlp_ratio",
    "softcap_logits",
    "final_softcap",
    "lm_head_precision_mode",
    "bos_token_id",
    "eos_token_id",
    "pad_token_id",
    "source_format",
    "checkpoint_step",
    "checkpoint_world_size",
)

HEAD_PRECISION_MODES: tuple[str, ...] = ("oracle_exact", "fp32_accumulate")

#: Values the modeling code implements. Anything else must fail loudly: every
#: entry here is a fork in the numerics, and a silent fallback would produce a
#: plausible-looking model with wrong numbers.
SUPPORTED = {
    "mlp_type": {"swiglu"},
    "pos_mode": {"rope"},
    "qk_norm": {"rms_pre_rope"},
    "rms_norm_eps_mode": {"torch_finfo_default"},
    "rope_style": {"interleaved_pairs_odd_lane_sign_flip"},
    "rope_cos_sin_dtype": {"bfloat16"},
    "mudd_accumulation": {"ordered_left_to_right"},
    "ve_head_slice": {"first_num_key_value_heads"},
    "lm_head_precision_mode": set(HEAD_PRECISION_MODES),
    "softcap_kind": {"sigmoid"},
}

MLP_FORMULAS = {
    "swiglu": (
        "(silu(gate_proj(h)) * up_proj(h)) @ down_proj, no clamp on either factor"
    ),
}

#: The reference applies `k >= q - sliding_window`, i.e. the window is inclusive
#: of the query token, so a local layer sees `sliding_window + 1` keys. Matching
#: the raw number against an exclusive-convention kernel silently drops the
#: oldest key on every local layer.
SLIDING_WINDOW_CONVENTION = (
    "k >= q - sliding_window, inclusive of the query token "
    "(span = sliding_window + 1 keys)"
)


def _derive_layer_types(num_hidden_layers: int, global_layers: list[int]) -> list[str]:
    global_layer_set = set(global_layers)
    return [
        "full_attention" if layer_idx in global_layer_set else "sliding_attention"
        for layer_idx in range(num_hidden_layers)
    ]


class LimiteConfig(PretrainedConfig):
    """Transformers configuration for Limite models."""

    model_type = MODEL_TYPE
    keys_to_ignore_at_inference = ["past_key_values"]

    @classmethod
    def from_dict(cls, config_dict: dict[str, Any], **kwargs: Any) -> "LimiteConfig":
        missing = [
            field
            for field in REQUIRED_CONFIG_FIELDS
            if (
                config_dict.get("torch_dtype", config_dict.get("dtype"))
                if field == "torch_dtype"
                else config_dict.get(field)
            )
            is None
        ]
        if missing:
            raise ValueError(
                f"Limite config is missing required fields {missing}. "
                "Re-export the checkpoint "
                "instead of guessing numerical choices."
            )
        return super().from_dict(config_dict, **kwargs)

    def __init__(
        self,
        vocab_size: int = 151680,
        padded_vocab_size: int | None = None,
        tokenizer_vocab_size: int | None = None,
        hidden_size: int = 1280,
        intermediate_size: int = 5120,
        num_hidden_layers: int = 48,
        num_attention_heads: int = 10,
        num_key_value_heads: int = 2,
        head_dim: int = 128,
        mlp_ratio: int = 4,
        mlp_type: str = "swiglu",
        mlp_formula: str | None = None,
        max_position_embeddings: int = 8192,
        tie_word_embeddings: bool = True,
        rms_norm_has_weight: bool = False,
        rms_norm_eps_mode: str = "torch_finfo_default",
        qk_norm: str = "rms_pre_rope",
        attention_softmax_scale: float = 0.1,
        sliding_window: int = 1024,
        sliding_window_convention: str = SLIDING_WINDOW_CONVENTION,
        global_window: int = -1,
        global_layers: list[int] | None = None,
        global_every: int = 4,
        global_nope: bool = True,
        attn_gate_channels: int = 0,
        attn_gate_scale: float = 2.0,
        attn_gate_applied: str = "per_head_before_o_proj",
        attention_dropout: float = 0.0,
        pos_mode: str = "rope",
        rope_frac: float = 0.5,
        rope_base_local: float = 1024.0,
        rope_base_global: float = 1024.0,
        rope_per_layer: bool = False,
        rope_n_pairs: int | None = None,
        rope_style: str = "interleaved_pairs_odd_lane_sign_flip",
        rope_cos_sin_dtype: str = "bfloat16",
        ve_dim: int = 128,
        ve_layers: list[int] | None = None,
        ve_gate_channels: int = 12,
        ve_gate_scale: float = 2.0,
        ve_head_slice: str = "first_num_key_value_heads",
        ve_stored_heads: int | None = None,
        ve_applied_before_qk_norm: bool = True,
        xsa: bool = True,
        xsa_layers: list[int] | None = None,
        xsa_normalize_eps: float = 1e-4,
        mudd: bool = True,
        mudd_at: list[int] | None = None,
        mudd_layers: list[int] | None = None,
        mudd_taps: int = 3,
        mudd_inter: int = 32,
        mudd_tap_idx: dict[str, list[int]] | None = None,
        mudd_hist_convention: str = (
            "hist[0] = rms_norm(embedding); hist[j] = output of layer j-1"
        ),
        mudd_accumulation: str = "ordered_left_to_right",
        mudd_mlp: bool = False,
        mudd_r_site: str | None = None,
        softcap_logits: dict[str, Any] | None = None,
        final_softcap: float = 0.0,
        lm_head_precision_mode: str = "oracle_exact",
        lm_head_compute_dtype: str | None = None,
        pad_token_id: int | None = PAD_TOKEN_ID,
        bos_token_id: int | None = BOS_TOKEN_ID,
        eos_token_id: int | list[int] | None = EOS_TOKEN_ID,
        source_format: str | None = None,
        source_metadata_assumptions: list[str] | None = None,
        checkpoint_step: int | None = None,
        checkpoint_world_size: int | None = None,
        **kwargs: Any,
    ):
        kwargs.setdefault("architectures", [ARCHITECTURE])
        kwargs.setdefault("attn_implementation", "sdpa")
        self.vocab_size = vocab_size
        self.padded_vocab_size = (
            padded_vocab_size if padded_vocab_size is not None else vocab_size
        )
        self.tokenizer_vocab_size = tokenizer_vocab_size
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.num_key_value_heads = num_key_value_heads
        self.head_dim = head_dim
        self.mlp_ratio = mlp_ratio
        self.mlp_type = mlp_type
        self.mlp_formula = (
            mlp_formula if mlp_formula is not None else MLP_FORMULAS.get(mlp_type)
        )
        self.max_position_embeddings = max_position_embeddings
        self.rms_norm_has_weight = rms_norm_has_weight
        self.rms_norm_eps_mode = rms_norm_eps_mode
        self.qk_norm = qk_norm
        self.attention_softmax_scale = attention_softmax_scale
        self._serialized_sliding_window = int(sliding_window)
        self.sliding_window = self._serialized_sliding_window + 1
        self.sliding_window_convention = sliding_window_convention
        self.global_window = global_window
        self.global_every = global_every
        self.global_layers = sorted(int(x) for x in (global_layers or []))
        self.global_nope = global_nope
        self.attn_gate_channels = attn_gate_channels
        self.attn_gate_scale = attn_gate_scale
        self.attn_gate_applied = attn_gate_applied
        self.attention_dropout = attention_dropout
        self.pos_mode = pos_mode
        self.rope_frac = rope_frac
        self.rope_base_local = rope_base_local
        self.rope_base_global = rope_base_global
        self.rope_per_layer = rope_per_layer
        self.rope_n_pairs = (
            rope_n_pairs
            if rope_n_pairs is not None
            else max(1, int(head_dim * rope_frac) // 2)
        )
        self.rope_style = rope_style
        self.rope_cos_sin_dtype = rope_cos_sin_dtype
        self.ve_dim = ve_dim
        self.ve_layers = sorted(int(x) for x in (ve_layers or []))
        self.ve_gate_channels = ve_gate_channels
        self.ve_gate_scale = ve_gate_scale
        self.ve_head_slice = ve_head_slice
        self.ve_stored_heads = (
            ve_stored_heads if ve_stored_heads is not None else num_attention_heads
        )
        self.ve_applied_before_qk_norm = ve_applied_before_qk_norm
        self.xsa = xsa
        self.xsa_layers = sorted(int(x) for x in (xsa_layers or []))
        self.xsa_normalize_eps = xsa_normalize_eps
        self.mudd = mudd
        self.mudd_at = sorted(int(x) for x in (mudd_at or []))
        self.mudd_layers = sorted(
            int(x) for x in (mudd_layers if mudd_layers is not None else self.mudd_at)
        )
        self.mudd_taps = mudd_taps
        self.mudd_inter = mudd_inter
        self.mudd_tap_idx = {
            str(k): [int(i) for i in v] for k, v in (mudd_tap_idx or {}).items()
        }
        self.mudd_hist_convention = mudd_hist_convention
        self.mudd_accumulation = mudd_accumulation
        self.mudd_mlp = mudd_mlp
        self.mudd_r_site = (mudd_r_site or "resid") if mudd_mlp else None
        self.softcap_logits = dict(
            softcap_logits or {"kind": "sigmoid", "a": 23.0, "b": 5.0, "c": 7.5}
        )
        self.final_softcap = final_softcap
        if lm_head_compute_dtype is not None:
            raise ValueError(
                "Limite config carries the superseded "
                f"lm_head_compute_dtype={lm_head_compute_dtype!r}. That field named "
                "only one of the head's three roundable stages and is not "
                "reinterpreted; re-export the checkpoint."
            )
        self.lm_head_precision_mode = lm_head_precision_mode
        self.source_format = source_format
        self.source_metadata_assumptions = source_metadata_assumptions
        self.checkpoint_step = checkpoint_step
        self.checkpoint_world_size = checkpoint_world_size
        super().__init__(
            pad_token_id=pad_token_id,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            tie_word_embeddings=tie_word_embeddings,
            **kwargs,
        )
        if getattr(self, "layer_types", None) is None:
            self.layer_types = _derive_layer_types(
                self.num_hidden_layers,
                self.global_layers,
            )
        self.validate_architecture()

    def to_dict(self) -> dict[str, Any]:
        """Serialize the checkpoint convention, not the native cache span."""
        output = super().to_dict()
        output["sliding_window"] = self._serialized_sliding_window
        output.pop("_serialized_sliding_window", None)
        return output

    @property
    def ve_layer_to_ordinal(self) -> dict[int, int]:
        return {layer: ordinal for ordinal, layer in enumerate(self.ve_layers)}

    def tap_indices(self, layer_idx: int) -> list[int] | None:
        return self.mudd_tap_idx.get(str(layer_idx))

    def is_global_layer(self, layer_idx: int) -> bool:
        return layer_idx in set(self.global_layers)

    def validate_architecture(self) -> None:
        for field, allowed in (
            ("mlp_type", SUPPORTED["mlp_type"]),
            ("pos_mode", SUPPORTED["pos_mode"]),
            ("qk_norm", SUPPORTED["qk_norm"]),
            ("rms_norm_eps_mode", SUPPORTED["rms_norm_eps_mode"]),
            ("rope_style", SUPPORTED["rope_style"]),
            ("rope_cos_sin_dtype", SUPPORTED["rope_cos_sin_dtype"]),
            ("mudd_accumulation", SUPPORTED["mudd_accumulation"]),
            ("ve_head_slice", SUPPORTED["ve_head_slice"]),
            ("lm_head_precision_mode", SUPPORTED["lm_head_precision_mode"]),
        ):
            value = getattr(self, field)
            if value not in allowed:
                raise NotImplementedError(
                    f"Limite does not implement {field}={value!r}; "
                    f"supported: {sorted(allowed)}"
                )
        if self.rms_norm_has_weight:
            raise NotImplementedError(
                "Limite RMS norm carries no learnable gain "
                "(rms_norm_has_weight must be False)."
            )
        if self.sliding_window_convention != SLIDING_WINDOW_CONVENTION:
            raise NotImplementedError(
                f"Unexpected sliding_window_convention "
                f"{self.sliding_window_convention!r}. The window span is an "
                "off-by-one trap; refusing to guess."
            )
        if self.softcap_logits.get("kind") not in SUPPORTED["softcap_kind"]:
            raise NotImplementedError(
                "Limite implements only sigmoid logit softcapping, got "
                f"{self.softcap_logits!r}"
            )
        if self.final_softcap:
            raise NotImplementedError(
                "Limite does not implement the pre-head tanh cap "
                "(final_softcap must be 0)."
            )
        if self.attn_gate_channels < 0:
            raise ValueError("attn_gate_channels must be non-negative.")
        if self.attn_gate_channels:
            if self.attn_gate_scale != 2.0:
                raise NotImplementedError(
                    "Limite implements the attention gate only as "
                    f"2 * sigmoid(...), got {self.attn_gate_scale}."
                )
            if self.attn_gate_applied != "per_head_before_o_proj":
                raise NotImplementedError(
                    "Limite does not implement "
                    f"attn_gate_applied={self.attn_gate_applied!r}."
                )
        if not self.global_nope:
            raise NotImplementedError(
                "Limite implements rotary-free global layers only "
                "(global_nope must be True)."
            )
        if self.rope_per_layer and self.rope_base_local != self.rope_base_global:
            raise NotImplementedError(
                "rope_per_layer with distinct bases is unreachable while global "
                "layers skip rotary entirely."
            )
        if self.num_attention_heads % self.num_key_value_heads != 0:
            raise ValueError(
                f"num_attention_heads ({self.num_attention_heads}) must be "
                f"divisible by num_key_value_heads ({self.num_key_value_heads})."
            )
        if self.num_attention_heads * self.head_dim != self.hidden_size:
            raise ValueError(
                "num_attention_heads * head_dim "
                f"({self.num_attention_heads * self.head_dim}) must equal "
                f"hidden_size ({self.hidden_size})."
            )
        if self.mlp_formula != MLP_FORMULAS[self.mlp_type]:
            raise NotImplementedError(
                f"Limite does not implement mlp_formula={self.mlp_formula!r} "
                f"for {self.mlp_type!r}."
            )
        if self.ve_dim > self.head_dim:
            raise ValueError(
                f"ve_dim ({self.ve_dim}) must not exceed head_dim ({self.head_dim})."
            )
        if self.ve_gate_channels > self.hidden_size:
            raise ValueError(
                f"ve_gate_channels ({self.ve_gate_channels}) must not exceed "
                "hidden_size."
            )
        if self.ve_stored_heads not in {
            self.num_attention_heads,
            self.num_key_value_heads,
        }:
            raise ValueError(
                f"ve_stored_heads ({self.ve_stored_heads}) must be query-head "
                f"width ({self.num_attention_heads}) or key-value-head width "
                f"({self.num_key_value_heads})."
            )
        if self.padded_vocab_size != self.vocab_size:
            raise ValueError(
                f"vocab_size ({self.vocab_size}) is the width of the embedding "
                "and head matrices and must equal padded_vocab_size "
                f"({self.padded_vocab_size})."
            )
        for name, layers in (
            ("global_layers", self.global_layers),
            ("ve_layers", self.ve_layers),
            ("xsa_layers", self.xsa_layers),
            ("mudd_layers", self.mudd_layers),
        ):
            if layers and (layers[0] < 0 or layers[-1] >= self.num_hidden_layers):
                raise ValueError(
                    f"{name}={layers} is out of range for "
                    f"num_hidden_layers={self.num_hidden_layers}."
                )
        if self.mudd:
            if sorted(int(k) for k in self.mudd_tap_idx) != list(self.mudd_layers):
                raise ValueError(
                    f"mudd_tap_idx keys {sorted(self.mudd_tap_idx)} must match "
                    f"mudd_layers {self.mudd_layers}."
                )
            for layer, taps in self.mudd_tap_idx.items():
                if len(taps) != self.mudd_taps:
                    raise ValueError(
                        f"mudd_tap_idx[{layer}] has {len(taps)} taps, expected "
                        f"{self.mudd_taps}."
                    )
                if max(taps) > int(layer):
                    raise ValueError(
                        f"mudd_tap_idx[{layer}]={taps} reads a history entry "
                        f"that does not exist yet at layer {layer}."
                    )
        elif self.mudd_layers:
            raise ValueError("mudd is disabled but mudd_layers is non-empty.")
        if self.mudd_mlp:
            if not self.mudd:
                raise ValueError("mudd_mlp requires the shared MUDD H-way mixer.")
            if self.mudd_r_site != "resid":
                raise NotImplementedError(
                    f"Limite implements MUDD R only at the residual base, got "
                    f"{self.mudd_r_site!r}."
                )
        if not self.xsa and self.xsa_layers:
            raise ValueError("xsa is disabled but xsa_layers is non-empty.")


__all__ = [
    "LimiteConfig",
    "MLP_FORMULAS",
    "REQUIRED_CONFIG_FIELDS",
    "SLIDING_WINDOW_CONVENTION",
]