# 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[] 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"]