configuration_agnes.py

# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from huggingface_hub.dataclasses import strict

from transformers.configuration_utils import PreTrainedConfig
from transformers.modeling_rope_utils import RopeParameters
from transformers.utils import auto_docstring


# The two MoE block kinds and the three attention block kinds Agnes ships, plus
# the legacy per-layer integer -> attention-kind map used when a checkpoint still
# carries the old `compress_ratios` list.
AGNES_MLP_LAYER_TYPES = ("agnes_hash_moe", "agnes_moe")

AGNES_LAYER_TYPES = (
    "agnes_local_attention",
    "agnes_sparse_attention",
    "agnes_pooled_attention",
)

_COMPRESS_RATIO_TO_LAYER_TYPE = {
    0: "agnes_local_attention",
    4: "agnes_sparse_attention",
    128: "agnes_pooled_attention",
}


@auto_docstring
@strict
class AgnesConfig(PreTrainedConfig):
    r"""
    scoring_func (`str`):
        Activation applied to the router logits — one of `sqrtsoftplus`,
        `softmax`, `sigmoid`.
    rope_theta (`float`):
        Rotary base for the main (local) attention path.
    layer_types (`list[str]`):
        The per-layer attention schedule, drawn from `agnes_local_attention`,
        `agnes_sparse_attention`, `agnes_pooled_attention`. Default schedule:
        two pooled layers to bootstrap, then an alternating sparse/pooled tail.
    compress_rates (`dict[str, int]`):
        Compression stride per attention type; default
        `{"agnes_sparse_attention": 4, "agnes_pooled_attention": 128}`. For
        backward compatibility, top-level `compress_rate_csa` /
        `compress_rate_hca` kwargs are absorbed in `__post_init__`.
    compress_rope_theta (`float`):
        Rotary base shared by the compressed branches (used together with
        `rope_scaling` for YaRN).
    hc_mult (`int`):
        Number of parallel residual streams carried by the manifold-constrained
        hyper-connections (always on).
    hc_sinkhorn_iters (`int`):
        Iteration budget for the Sinkhorn-Knopp projection of the stream-mixing
        matrix onto the doubly-stochastic manifold.
    hc_eps (`float`):
        Small constant guarding the Sinkhorn normalisation against divide-by-zero.
    mlp_layer_types (`list[str]`):
        The per-layer MoE schedule over `agnes_hash_moe` / `agnes_moe`. Hash
        layers pick experts from the frozen `tid2eid[input_ids]` table; plain
        layers use learned top-k routing. Default: the first three layers hash,
        the rest learned. Legacy `num_hash_layers` is folded in at
        `__post_init__`.
    swiglu_limit (`float`):
        Clamp bound on the routed experts' gate/up pre-activations.
    parallel_ffn_intermediate_size (`int`):
        Width of the optional parallel dense-FFN branch on every non-hash MoE
        layer (`0` turns it off). Its `down_proj` is exported as zeros, so the
        branch is a no-op until trained.
    sliding_window (`int`):
        Local attention window width used by every block.
    o_groups (`int`):
        Head-group count for the grouped output projection.
    o_lora_rank (`int`):
        Per-group bottleneck width of the grouped output projection.
    index_n_heads (`int`):
        Query-head count of the Lightning Indexer.
    index_head_dim (`int`):
        Per-head width inside the Lightning Indexer.
    index_topk (`int`):
        How many compressed entries the indexer keeps per query.
    num_nextn_predict_layers (`int`):
        Number of MTP layers present in the source checkpoint (not built here).
    partial_rotary_factor (`float`, *optional*):
        Fraction of each head that is rotated. Defaults to
        `qk_rope_head_dim / head_dim`, sizing cos/sin to `qk_rope_head_dim`.
    """

    model_type = "agnes"
    keys_to_ignore_at_inference = ["past_key_values"]

    # --- core dimensions ---
    vocab_size: int = 129280
    hidden_size: int = 4096
    num_hidden_layers: int = 43

    # --- self-attention (shared-KV MQA + grouped output projection) ---
    num_attention_heads: int = 64
    num_key_value_heads: int = 1
    head_dim: int = 512
    default_partial_rotary_factor = 64 / 512  # `qk_rope_head_dim` (64) / `head_dim` (512)
    partial_rotary_factor: float | None = None
    q_lora_rank: int = 1024
    o_groups: int = 8
    o_lora_rank: int = 1024
    sliding_window: int = 128
    attention_bias: bool = False
    attention_dropout: float = 0.0

    # --- long-range compressor + lightning indexer ---
    index_n_heads: int = 64
    index_head_dim: int = 128
    index_topk: int = 512
    layer_types: list[str] | None = None
    compress_rates: dict | None = None
    default_compress_rates = {"agnes_sparse_attention": 4, "agnes_pooled_attention": 128}
    compress_rope_theta: float | int = 160000.0

    # --- mixture-of-experts routing ---
    moe_intermediate_size: int = 2048
    parallel_ffn_intermediate_size: int = 0
    n_routed_experts: int = 256
    n_shared_experts: int = 1
    num_experts_per_tok: int = 6
    mlp_layer_types: list[str] | None = None
    default_num_hash_layers = 3
    scoring_func: str = "sqrtsoftplus"
    norm_topk_prob: bool = True
    routed_scaling_factor: float = 1.5
    num_nextn_predict_layers: int = 1
    output_router_logits: bool = False
    router_aux_loss_coef: float = 0.001
    router_jitter_noise: float = 0.0

    # --- feed-forward activation ---
    hidden_act: str = "silu"
    swiglu_limit: float = 10.0
    mlp_bias: bool = False

    # --- manifold-constrained hyper-connections ---
    hc_mult: int = 4
    hc_sinkhorn_iters: int = 20
    hc_eps: float = 1.0e-6

    # --- normalisation + initialisation ---
    rms_norm_eps: float = 1.0e-6
    initializer_range: float = 0.02

    # --- rotary position embedding ---
    rope_theta: float | int = 10000.0
    rope_parameters: RopeParameters | dict | None = None
    max_position_embeddings: int = 1048576

    # --- special tokens + weight tying ---
    pad_token_id: int | None = None
    bos_token_id: int | None = 0
    eos_token_id: int | list[int] | None = 1
    tie_word_embeddings: bool = False

    # --- runtime ---
    use_cache: bool = True

    # ------------------------------------------------------------------ #
    # Non-field class attributes: alias map, parallelism plans, rope labels
    # ------------------------------------------------------------------ #

    # Expert-parallel plan. Agnes ships EP only — it is MoE, so there is no
    # `base_model_tp_plan`. The gate routes, the routed experts run as a grouped
    # GEMM sharded on the expert axis, and the experts module is wrapped so its
    # output is all-reduced. Core attention stays replicated: it is shared-KV MQA
    # broadcasting a single KV head to every query head, so colwise-sharding
    # `q_b_proj` would desync the KV broadcast from the rank-local head count; the
    # small shared MLP is not worth sharding either. The one exception is the
    # Lightning Indexer, whose keys are replicated (its own compressor runs at
    # index_head_dim on replicated hidden states): there `q_b_proj` and
    # `scorer.weights_proj` go colwise and the `scorer` output is all-reduced so
    # every rank picks the same top-k.
    base_model_ep_plan = {
        "layers.*.mlp.gate": "ep_router",
        "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
        "layers.*.mlp.experts.down_proj": "grouped_gemm",
        "layers.*.mlp.experts": "moe_tp_experts",
        "layers.*.self_attn.compressor.indexer.q_b_proj": "colwise",
        "layers.*.self_attn.compressor.indexer.scorer.weights_proj": "colwise",
        "layers.*.self_attn.compressor.indexer.scorer": "all_reduce",
    }

    base_model_pp_plan = {
        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
        "norm": (["hidden_states"], ["hidden_states"]),
    }

    # `num_local_experts` / `intermediate_size` are the names the MoE FP8 / TP
    # integrations and the shared-expert MLP read; Agnes only stores
    # `n_routed_experts` / `moe_intermediate_size`, so alias them through.
    attribute_map = {
        "num_local_experts": "n_routed_experts",
        "intermediate_size": "moe_intermediate_size",
    }

    # `rope_parameters` is keyed by rope label (`main` / `compress`), not by
    # `layer_types`. The base `validate_rope` assumes the latter, so it is
    # overridden below to walk the rope-label sub-dicts directly.
    _rope_type_labels = ("main", "compress")

    def __post_init__(self, **kwargs):
        # Older checkpoints still ship a handful of legacy kwargs; pop them out of
        # the dict before the strict parent init runs, then reconcile each into its
        # current field afterwards.
        leg_ratios = kwargs.pop("compress_ratios", None)
        leg_csa = kwargs.pop("compress_rate_csa", None)
        leg_hca = kwargs.pop("compress_rate_hca", None)
        leg_hash = kwargs.pop("num_hash_layers", None)
        leg_qk_rope = kwargs.pop("qk_rope_head_dim", None)
        PreTrainedConfig.__post_init__(self, **kwargs)
        n = self.num_hidden_layers

        # compress_rates: seed the per-type defaults, then let a legacy scalar win.
        if self.compress_rates is None:
            self.compress_rates = dict(self.default_compress_rates)
        if leg_csa is not None:
            self.compress_rates["agnes_sparse_attention"] = leg_csa
        if leg_hca is not None:
            self.compress_rates["agnes_pooled_attention"] = leg_hca

        # layer_types precedence: explicit field, then legacy 0/4/128 ratios, then
        # the default (two pooled bootstrap layers + sparse/pooled interleave).
        if self.layer_types is None and leg_ratios is not None:
            self.layer_types = [_COMPRESS_RATIO_TO_LAYER_TYPE[r] for r in leg_ratios]
        if self.layer_types is None:
            tail = [
                "agnes_sparse_attention" if i % 2 else "agnes_pooled_attention"
                for i in range(max(n - 2, 0))
            ]
            self.layer_types = ["agnes_pooled_attention"] * min(n, 2) + tail
        self.layer_types = list(self.layer_types[:n])

        # mlp_layer_types: leading hash-routed layers, learned routing for the rest.
        if self.mlp_layer_types is None:
            n_hash = leg_hash if leg_hash is not None else self.default_num_hash_layers
            self.mlp_layer_types = ["agnes_hash_moe"] * min(n, n_hash) + ["agnes_moe"] * max(0, n - n_hash)
        self.mlp_layer_types = list(self.mlp_layer_types[:n])

        # partial_rotary_factor from a legacy qk_rope_head_dim if provided, else the
        # default; qk_rope_head_dim itself is only ever a runtime attr, not a field.
        if self.partial_rotary_factor is None:
            self.partial_rotary_factor = (
                leg_qk_rope / self.head_dim if leg_qk_rope is not None else self.default_partial_rotary_factor
            )
        self.qk_rope_head_dim = int(self.head_dim * self.partial_rotary_factor)

        # Normalise rope_parameters into the {main, compress} nesting. Local layers
        # use plain rope at rope_theta; only the compress branch may be YaRN, and if
        # it is we pin attention_factor=1.0 (Agnes never applies YaRN's mscale).
        rp = self.rope_parameters or {}
        if isinstance(rp.get("main"), dict) and isinstance(rp.get("compress"), dict):
            self.rope_parameters = {"main": rp["main"], "compress": rp["compress"]}
        else:
            extra = {k: v for k, v in rp.items() if k not in ("main", "compress")}
            main = {
                "rope_type": "default",
                "rope_theta": self.rope_theta,
                "partial_rotary_factor": self.partial_rotary_factor,
            }
            compress = {
                **extra,
                "rope_theta": self.compress_rope_theta,
                "partial_rotary_factor": self.partial_rotary_factor,
            }
            compress.setdefault("rope_type", "default")
            if compress["rope_type"] == "yarn":
                compress.setdefault("attention_factor", 1.0)
            self.rope_parameters = {"main": main, "compress": compress}

    def validate_layer_type(self):
        """Keep `layer_types` / `mlp_layer_types` within the block kinds Agnes
        actually builds, alongside the usual length check against
        `num_hidden_layers`."""
        if self.num_hidden_layers is None:
            return
        schedules = (
            ("layer_types", self.layer_types, AGNES_LAYER_TYPES),
            ("mlp_layer_types", self.mlp_layer_types, AGNES_MLP_LAYER_TYPES),
        )
        for name, schedule, allowed in schedules:
            if schedule is None:
                continue
            if len(schedule) != self.num_hidden_layers:
                raise ValueError(
                    f"`num_hidden_layers` ({self.num_hidden_layers}) must equal `len({name})` ({len(schedule)})."
                )
            unknown = [t for t in schedule if t not in allowed]
            if unknown:
                raise ValueError(f"`{name}` entries must be one of {allowed} for Agnes; got {unknown}.")

    def validate_rope(self):
        # The stock validators index self.rope_parameters[<key>] directly, which
        # breaks against the {main, compress} nesting, so run each rope-label
        # sub-dict through its own validator with self.rope_parameters temporarily
        # pointed at it, restoring the nested dict afterwards.
        nested = getattr(self, "rope_parameters", None) or {}
        ignore_keys = self.ignore_keys_at_rope_validation
        for label in self._rope_type_labels:
            sub = nested.get(label)
            if not isinstance(sub, dict):
                continue
            kind = sub.get("rope_type", sub.get("type", "default"))
            sub["rope_type"] = kind
            validator = getattr(self, f"_validate_{kind}_rope_parameters", None)
            if validator is None:
                continue
            self.rope_parameters = sub
            try:
                validator(sub, ignore_keys=ignore_keys)
            finally:
                self.rope_parameters = nested


__all__ = ["AgnesConfig"]