# 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. """ Laguna configuration for transformers 4.56-4.x (used by vLLM). This uses rope_theta + rope_scaling (legacy format) instead of rope_parameters (v5 format). """ from transformers.configuration_utils import PretrainedConfig 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`, *optional*, defaults to `True`): Whether to use softplus output gating on attention. When True, a g_proj linear layer is added and attn_output = attn_output * softplus(g_proj(x)). 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. rope_theta (`float`, *optional*, defaults to 500000.0): Base frequency for RoPE embeddings. rope_scaling (`dict`, *optional*): RoPE scaling configuration (e.g. YaRN, linear). 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 ``"global_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. 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. """ model_type = "laguna" keys_to_ignore_at_inference = ["past_key_values"] 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 = 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_theta: float = 500000.0, rope_scaling: dict | None = None, attention_dropout: float = 0.0, sliding_window: int | None = None, layer_types: list[str] | None = None, swa_attention_sink_enabled: bool = False, 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, 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] 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_theta = rope_theta self.rope_scaling = rope_scaling self.attention_dropout = attention_dropout # Sliding window attention arguments self.sliding_window = sliding_window self.layer_types = layer_types self.swa_attention_sink_enabled = swa_attention_sink_enabled # 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.output_router_logits = output_router_logits super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) __all__ = ["LagunaConfig"]