# ruff: noqa # Copyright 2025 Poolside and 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 transformers.configuration_utils import PretrainedConfig as PreTrainedConfig try: from transformers.modeling_rope_utils import RopeParameters except ImportError: RopeParameters = dict try: from transformers.utils.import_utils import is_causal_conv1d_available, is_flash_linear_attention_available except ImportError: def is_causal_conv1d_available(): return False def is_flash_linear_attention_available(): return False class LagunaConfig(PreTrainedConfig): r""" Configuration class for Laguna model. Laguna is Poolside's MoE architecture with: - Attention output gating (softplus gate) - Sigmoid routing instead of softmax - No QKV bias - Explicit head_dim parameter Args: head_dim (`int`, *optional*, defaults to 128): Dimension of attention heads. Laguna uses explicit head_dim rather than computing it from hidden_size // num_attention_heads. qkv_bias (`bool`, *optional*, defaults to `False`): Whether to add bias to QKV projections. Laguna uses no QKV bias. attention_bias (`bool`, *optional*, defaults to `False`): Whether to add bias to attention output projection. Laguna uses no attention bias. gating (`bool` or `str`, *optional*, defaults to `True`): Attention output gating mode. When ``True`` or ``"per-element"`` a g_proj linear layer with output size ``num_attention_heads * head_dim`` is added and ``attn_output = attn_output * softplus(g_proj(x))``. When ``"per-head"`` g_proj has output size ``num_attention_heads`` and the gate broadcasts across ``head_dim``. When ``False`` no gating is applied. partial_rotary_factor (`float`, *optional*): Fraction of head_dim to apply rotary embeddings to. When set, this value is injected into ``rope_parameters`` (and ``swa_rope_parameters``) if not already specified there. When ``None`` the default behaviour of the rope implementation is used (typically full rotary). num_attention_heads_per_layer (`list[int]`, *optional*): Optional per-layer override for ``num_attention_heads``. When provided the list length must equal ``num_hidden_layers`` and each entry is the head count used by that layer. When ``None`` every layer uses ``num_attention_heads``. vocab_size (`int`, *optional*, defaults to 100352): Vocabulary size of the Laguna model. hidden_size (`int`, *optional*, defaults to 2048): Dimension of the hidden representations. intermediate_size (`int`, *optional*, defaults to 8192): Dimension of the MLP representations for dense layers. num_hidden_layers (`int`, *optional*, defaults to 48): Number of hidden layers in the Transformer. num_attention_heads (`int`, *optional*, defaults to 32): Number of attention heads. num_key_value_heads (`int`, *optional*, defaults to 8): Number of key-value heads for GQA. max_position_embeddings (`int`, *optional*, defaults to 4096): Maximum sequence length. rms_norm_eps (`float`, *optional*, defaults to 1e-6): Epsilon for RMSNorm layers. sliding_window (`int`, *optional*): Sliding window attention size. Used by layers whose type in ``layer_types`` is ``"sliding_attention"``. When ``None``, all layers use full attention. layer_types (`list[str]`, *optional*): Per-layer attention type. Each element should be ``"sliding_attention"`` or ``"full_attention"``. Length must equal ``num_hidden_layers``. When ``None``, all layers default to global attention. swa_attention_sink_enabled (`bool`, *optional*, defaults to `False`): Whether to enable learnable attention sinks on sliding-window attention layers. When enabled, a per-head bias parameter is added that allows the model to attend to position 0 even when it falls outside the sliding window. swa_rope_parameters (`RopeParameters`, *optional*): Separate RoPE configuration for sliding-window attention layers. When ``None``, SWA layers use the same RoPE as global attention layers. num_experts (`int`, *optional*, defaults to 256): Number of routed experts. num_experts_per_tok (`int`, *optional*, defaults to 16): Number of experts selected per token (top-k). moe_intermediate_size (`int`, *optional*, defaults to 1024): Intermediate size of routed experts. shared_expert_intermediate_size (`int`, *optional*, defaults to 1024): Intermediate size of the shared expert. norm_topk_prob (`bool`, *optional*, defaults to `True`): Whether to normalize top-k routing probabilities. decoder_sparse_step (`int`, *optional*, defaults to 1): Frequency of MoE layers (1 = every layer is MoE after mlp_only_layers). mlp_only_layers (`list[int]`, *optional*, defaults to `[0]`): Layer indices that use dense MLP instead of MoE. router_aux_loss_coef (`float`, *optional*, defaults to 0.001): Auxiliary loss coefficient for load balancing. moe_routed_scaling_factor (`float`, *optional*, defaults to 1.0): Scalar multiplier applied to the routed-expert output before combining with the shared-expert output. moe_apply_router_weight_on_input (`bool`, *optional*, defaults to `False`): When ``True`` the top-k routing weights are multiplied into each expert's input rather than its output. Matches the numerical form used by the trained checkpoint. moe_router_logit_softcapping (`float`, *optional*, defaults to 0.0): Optional soft-capping value ``c`` applied to router logits as ``x = tanh(x / c) * c`` before sigmoid + top-k. Disabled when ``0``. rope_parameters (`RopeParameters`, *optional*): RoPE configuration. Defaults to rope_theta=500000.0. """ model_type = "laguna" keys_to_ignore_at_inference = ["past_key_values"] # PreTrainedConfig in transformers v5 no longer auto-declares these; subclasses # opt in by providing class-level annotations with defaults. pad_token_id: int | None = None bos_token_id: int | None = None eos_token_id: int | list[int] | None = None base_model_tp_plan = { "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", "layers.*.self_attn.g_proj": "colwise", # Laguna-specific gating projection "layers.*.self_attn.o_proj": "rowwise", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", "layers.*.mlp.down_proj": "rowwise", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } def __init__( self, vocab_size: int = 100352, hidden_size: int = 2048, intermediate_size: int = 8192, num_hidden_layers: int = 48, num_attention_heads: int = 32, num_key_value_heads: int = 8, head_dim: int = 128, qkv_bias: bool = False, attention_bias: bool = False, gating: bool | str = True, hidden_act: str = "silu", max_position_embeddings: int = 4096, initializer_range: float = 0.02, rms_norm_eps: float = 1e-6, use_cache: bool = True, tie_word_embeddings: bool = False, rope_parameters: RopeParameters | dict[str, RopeParameters] | None = None, partial_rotary_factor: float | None = None, attention_dropout: float = 0.0, sliding_window: int | None = None, layer_types: list[str] | None = None, num_attention_heads_per_layer: list[int] | None = None, swa_attention_sink_enabled: bool = False, swa_rope_parameters: RopeParameters | None = None, num_experts: int = 256, num_experts_per_tok: int = 16, moe_intermediate_size: int = 1024, shared_expert_intermediate_size: int = 1024, norm_topk_prob: bool = True, decoder_sparse_step: int = 1, mlp_only_layers: list[int] | None = None, router_aux_loss_coef: float = 0.001, moe_routed_scaling_factor: float = 1.0, moe_apply_router_weight_on_input: bool = False, moe_router_logit_softcapping: float = 0.0, output_router_logits: bool = False, **kwargs, ): # Default mlp_only_layers: first layer is dense (moe_first_k_dense_replace=1) if mlp_only_layers is None: mlp_only_layers = [0] # Default layer_types: all layers use full attention (Laguna-M). Laguna-XS # ships an explicit list with a mix of "full_attention" and "sliding_attention". # Downstream mask builders (``create_masks_for_generate``) iterate # ``layer_types``, so it must be a list — not left as ``None``. if layer_types is None: layer_types = ["full_attention"] * num_hidden_layers # Default rope_parameters with Laguna's theta if rope_parameters is None: rope_parameters = {"rope_type": "default", "rope_theta": 500000.0} # config.json stores SWA rope nested in rope_parameters["sliding_attention"] # and carries no top-level swa_rope_parameters. Derive it here, else the # sliding-window layers silently reuse the full-attention rope. if swa_rope_parameters is None and isinstance(rope_parameters, dict): swa_rope_parameters = rope_parameters.get("sliding_attention") # If ``partial_rotary_factor`` is set at the top level, inject it into any # rope dict that does not already carry one so the rotary embedding picks # it up consistently for both full-attention and SWA layers. if partial_rotary_factor is not None: if isinstance(rope_parameters, dict) and "partial_rotary_factor" not in rope_parameters: rope_parameters = {**rope_parameters, "partial_rotary_factor": partial_rotary_factor} if isinstance(swa_rope_parameters, dict) and "partial_rotary_factor" not in swa_rope_parameters: swa_rope_parameters = { **swa_rope_parameters, "partial_rotary_factor": partial_rotary_factor, } self.vocab_size = 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.qkv_bias = qkv_bias self.attention_bias = attention_bias self.gating = gating self.hidden_act = hidden_act self.max_position_embeddings = max_position_embeddings self.initializer_range = initializer_range self.rms_norm_eps = rms_norm_eps self.use_cache = use_cache self.rope_parameters = rope_parameters self.partial_rotary_factor = partial_rotary_factor self.attention_dropout = attention_dropout # Sliding window attention arguments self.sliding_window = sliding_window self.layer_types = layer_types self.num_attention_heads_per_layer = num_attention_heads_per_layer self.swa_attention_sink_enabled = swa_attention_sink_enabled self.swa_rope_parameters = swa_rope_parameters # MoE arguments self.num_experts = num_experts self.num_experts_per_tok = num_experts_per_tok self.moe_intermediate_size = moe_intermediate_size self.shared_expert_intermediate_size = shared_expert_intermediate_size self.norm_topk_prob = norm_topk_prob self.decoder_sparse_step = decoder_sparse_step self.mlp_only_layers = mlp_only_layers self.router_aux_loss_coef = router_aux_loss_coef self.moe_routed_scaling_factor = moe_routed_scaling_factor self.moe_apply_router_weight_on_input = moe_apply_router_weight_on_input self.moe_router_logit_softcapping = moe_router_logit_softcapping self.output_router_logits = output_router_logits super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) __all__ = ["LagunaConfig"]