model_args.py

import dataclasses
from dataclasses import dataclass
from enum import Enum
from functools import partial
from typing import Any, Dict, Literal, Optional

import torch.nn.functional as F

MoEScoreFunc = Literal["sigmoid", "softmax", "softmax_all"]


class ActFn(Enum):
    SILU = "silu"
    GELU = "gelu"
    GELU_APPROX = "gelu_approx"

    @classmethod
    def from_string(cls, value: str) -> "ActFn":
        """Convert string to ActFn enum."""
        try:
            return cls(value)
        except ValueError:
            valid_values = [e.value for e in cls]
            raise ValueError(
                f"Invalid activation function: {value}. Valid options: {valid_values}"
            )

    def get_function(self):
        """Return the corresponding activation function."""
        if self == ActFn.SILU:
            return F.silu
        elif self == ActFn.GELU:
            return F.gelu
        elif self == ActFn.GELU_APPROX:
            return partial(F.gelu, approximate="tanh")
        else:
            raise ValueError(f"Unsupported activation function: {self}")


@dataclass
class ModelArgs:
    dim: int = 4096
    n_layers: int = 1
    n_heads: int = 32
    n_kv_heads: Optional[int] = None
    vocab_size: int = 512  # Arbitrary value, should be defined later by tokenizer.
    hidden_dim: Optional[int] = None
    head_dim: Optional[int] = None  # Optional customized head_dim
    # Qwen3.5 linear-attention dimensions.
    linear_conv_kernel_dim: int = 4
    linear_key_head_dim: Optional[int] = None
    linear_value_head_dim: Optional[int] = None
    linear_num_key_heads: Optional[int] = None
    linear_num_value_heads: Optional[int] = None
    # Qwen3.5 RMSNorm uses (1 + weight) scaling.
    rms_norm_add_unit_offset: bool = False
    multiple_of: int = 256  # make SwiGLU hidden layer size multiple of large power of 2
    ffn_dim_multiplier: Optional[float] = None
    model_architecture: str = (
        "LlamaForCausalLM"  # This setting is currently only supported for the QNN backend
    )
    attention_multiplier: Optional[float] = (
        None  # Scaling factor 1/sqrt(d_k) in attention formula
    )
    norm_eps: float = 1e-5
    post_attention_norm: bool = False
    post_ffn_norm: bool = False
    max_batch_size: int = 1
    max_seq_len: int = 2048
    max_context_len: int = 2048
    use_ffn_norm: bool = True
    output_bias: bool = False
    moe: bool = False  # True to enable the MoE (Mixture of Experts)
    num_experts: int = 8  # Number of experts
    num_activated_experts: int = 2  # Number of experts to activate
    moe_hidden_dim: Optional[int] = None
    moe_shared_expert_hidden_dim: Optional[int] = None
    moe_use_expert_bias: bool = False
    moe_score_func: MoEScoreFunc = "softmax"
    moe_route_scale: float = 1.0
    attention_type: str = "mha"  # Attention type, registered in attention.py
    use_q_gate: bool = (
        False  # Use q-gated projection in attention (Qwen3.5 full attention)
    )
    norm_type: str = "rmsnorm"  # Normalization type, registered in norm.py
    act_fn: ActFn = dataclasses.field(default=ActFn.SILU)  # Activation function type
    attention_qkv_bias: bool = False
    use_kv_cache: bool = False  # Use key/value cache
    use_sdpa_with_kv_cache_op: bool = (
        False  # Use custom sdpa op that updates kv cache in-place
    )
    # Device to use for the model: "cpu" or "cuda" (needed for QAT)
    # Only used for creating Rope parameters
    device: str = "cpu"
    # Generate logits for all inputs. When it's True, it would take big memory usage
    # at runtime. Enable it only necessary (e.g., use perplexity tools that requires
    # logits for all input tokens.)
    generate_full_logits: bool = False
    logits_scaling: Optional[float] = (
        None  # Scaling factor applied to the logits of model, functioning similarly to a temperature parameter.
    )
    enable_dynamic_shape: bool = False  # export model with dynamic shape support
    # A dictionary mapping from pruned token-id to original token-id
    input_prune_map: Optional[Dict[int, int]] = None
    # A dictionary mapping from pruned token-id to original token-id
    output_prune_map: Optional[Dict[int, int]] = None
    apply_embedding: bool = True  # Use embedding inside the transformer
    embedding_scale_factor: float = 1.0  # Multiple by which to scale embeddings.
    apply_output: bool = True  # Use output layer (unembedding) inside the transformer
    use_qk_norm: bool = False  # apply normalization to q and k in the attention
    qk_norm_before_rope: bool = False  # when to apply qk norm
    qk_norm_affine: bool = (
        True  # whether QK norm has learnable weight (False = scaleless)
    )
    residual_multiplier: Optional[float] = (
        None  # Scaling factor applied to the residual hidden states
    )
    use_hf_rope: bool = False  # Use HuggingFace's RoPE implementation
    no_rope_layer_interval: Optional[int] = (
        None  # Interval at which to skip RoPE. From Rope to Nope and Back Again: A New Hybrid Attention Strategy (https://huggingface.co/papers/2501.18795).
    )
    partial_rotary_factor: float = 1.0
    rope_theta: Optional[float] = (
        None  # The official name to override self.rope_freq_base.
    )
    local_rope_theta: Optional[float] = (
        None  # For sliding window attention. e.g., gemma3-1b
    )
    rope_parameters: Optional[Dict[str, Dict[str, Any]]] = (
        None  # Per-layer-type RoPE configs. e.g., {"full_attention": {"rope_theta": 5000000, "partial_rotary_factor": 0.25}}
    )
    rope_freq_base: float = 10000.0  # The base frequency for RoPE. Keep it for BC.
    use_scaled_rope: bool = False  # Use scaled RoPE, introduced in llama3.1.
    rope_scale_factor: int = 8
    high_freq_factor: int = 4
    # LongRoPE (https://arxiv.org/abs/2402.13753) used by Phi-3 / Phi-4 family.
    # Mirrors HF's rope_scaling.{short_factor,long_factor,attention_factor}
    # plus original_max_position_embeddings / max_position_embeddings.
    rope_scaling_short_factor: Optional[list] = None
    rope_scaling_long_factor: Optional[list] = None
    original_max_position_embeddings: Optional[int] = None
    max_position_embeddings: Optional[int] = None
    rope_scaling_attention_factor: Optional[float] = None
    # Additional Model Metadata needed at runtime
    bos_idx: int = 1
    eos_idx: int = 3
    bos_count: int = -1  # i.e., a single EOS is used as BOS
    eos_count: int = 2

    quantization_args: Optional[dict] = None
    # LoRA for QAT.
    lora_args: Optional[dict] = None

    # LoRA arguments to set up a LoRA inference model.
    # These arguments come directly from a torchtune adapter_config.json file.
    r: Optional[int] = None  # Rank.
    lora_alpha: Optional[int] = None  # Alpha.
    # Modules that we can apply lora adapters to.
    # Eg. q_proj, k_proj, v_proj, output_proj/o_proj, down_proj, gate_proj, up_proj
    target_modules: Optional[list] = None
    peft_type: Optional[str] = None  # PEFT type.
    base_model_name_or_path: Optional[str] = None  # Base model name or path.
    kv_io_bit_width: Optional[int] = (
        None  # KV cache bit width. This is for QNN backend only for now.
    )
    attention_kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict)
    # Hybrid models can have layer types different from attention
    layer_types: Optional[list] = None
    # Per-layer MLP type: "default" for standard FFN, "skip" for no FFN block.
    # Indexed by layer id (e.g. mlp_type[0] applies to layer 0).
    mlp_type: Optional[list] = None
    model_architecture: Optional[str] = (
        None  # Architecture of model. For HF models, please refer to the HF model.config.architectures. This is used in QNN backend only for now.
    )
    sliding_window: Optional[int] = (
        None  # sliding window size for sliding window attention
    )
    # YOCO KV sharing: number of layers that share KV cache
    # with earlier donor layers instead of computing their own K/V projections.
    # When > 0, the last num_kv_shared_layers share KV from earlier layers.
    num_kv_shared_layers: int = 0
    # gemma2 attn and output soft capping
    final_logit_softcapping: Optional[float] = None
    attn_logit_softcapping: Optional[float] = None

    # rlformers forward-pass features for on-device model parity
    normalize_tok_embeddings: bool = False
    scale_query_by: float = 1.0
    use_attn_o_gate: bool = False
    headwise_attn_output_gate: bool = False
    use_attn_o_norm: bool = False
    use_residual_gate: bool = False
    use_ffn_learnable_scales: bool = False
    # Zero-centered gamma: the checkpoint stores gamma offset by -1, so the
    # effective scale is ``weight + 1``. rlformers applies this in
    # RMSNormWithInputScale; ignoring it silently rescales every post-FFN norm.
    norm_zero_centered_gamma: bool = False
    output_norm_gain_center_type: str = "one"
    output_soft_cap_temp: Optional[float] = None
    output_linear_intermediate_dim: int = 0

    def __post_init__(self):  # noqa: C901
        if self.n_kv_heads is None:
            self.n_kv_heads = self.n_heads

        # rope_theta overrides rope_freq_base since it's the official name.
        if self.rope_theta is not None:
            self.rope_freq_base = self.rope_theta

        if self.use_sdpa_with_kv_cache_op:
            assert self.use_kv_cache, "use_sdpa_with_kv_cache_op requires use_kv_cache"

        if self.hidden_dim is None:
            # If hidden_dim is not explicitly set in the ModelArgs,
            # then calculate implicitly based on dim and also multiple of `args.multiple_of`
            multiple_of = self.multiple_of
            hidden_dim = 4 * self.dim
            hidden_dim = int(2 * hidden_dim / 3)
            if self.ffn_dim_multiplier is not None:
                hidden_dim = int(self.ffn_dim_multiplier * hidden_dim)

            def find_multiple(n: int, k: int) -> int:
                if n % k == 0:
                    return n
                return n + k - (n % k)

            self.hidden_dim = find_multiple(hidden_dim, multiple_of)

        if self.head_dim is None:
            self.head_dim = self.dim // self.n_heads

        if self.linear_key_head_dim is None:
            self.linear_key_head_dim = self.head_dim
        if self.linear_value_head_dim is None:
            self.linear_value_head_dim = self.head_dim
        if self.linear_num_key_heads is None:
            self.linear_num_key_heads = self.n_heads
        if self.linear_num_value_heads is None:
            self.linear_num_value_heads = self.n_heads

        # Convert string act_fn to enum if needed
        if isinstance(self.act_fn, str):
            self.act_fn = ActFn.from_string(self.act_fn)