Azrail/smallm_70

1import logging2from transformers import PretrainedConfig3from typing import Optional4 5logger = logging.getLogger(__name__)6 7 8class SmalLmConfig(PretrainedConfig):9    """10    Base config for all SmalLm models11 12    Raises:13        ValueError: Positional_bias_type must be in suported types14        ValueError: In case of rope positional_bias_type head_size can't be anything15    """16    model_type = "smallm"17 18    def __init__(19        self,20        # global model params21        hidden_size: int = 512,22        intermediate_size: int = 2048,23        mlp_bias: bool = False,24        num_hidden_layers: int = 27,25        rms_norm_eps: float = 1e-6,26        rms_affine: bool = False,27        initializer_range: float = 0.02,28        output_hidden_states: bool = False,29        output_attentions: bool = False,30        use_cache: bool = True,31        sliding_window_attention: bool = True,32        sliding_window_context: int = 1024,33        sliding_window_period: int = 4,34        embedding_dropout: float = 0.0,35        layer_dropout: float = 0.1,36        max_seq_len: int = 2048,37        original_seq_len: int | None = None,38        tie_word_embeddings: bool = True,39        # attention params40        num_attention_heads: int = 9,41        num_kv_heads: int = 3,42        head_size: Optional[int] = None,43        attention_dropout: float = 0.1,44        positional_bias_type: str = "rope",45        high_rotations: int = 32,46        low_rotations: int = 1,47        attention_bias: bool = False,48        rope_base: int = 100000,49        # MoE params50        use_moe: bool = True,51        moe_period: int = 3,52        expert_size: int = 256,53        shared_experts: int = 2,54        routed_experts: int = 16,55        token_experts: int = 4,56        noisy_experts: bool = False,57        moe_bias: bool = False,58        balancing_coef: float = 1e-4,59        no_moe_layers: int = 5,60        # extra params61        vocab_size: int = 60000,62        bos_token_id: int = 1,63        eos_token_id: int = 0,64        pad_token_id: int = 0,65        static_residual: bool = False,66        **kwargs,67    ):68        if positional_bias_type not in ["alibi", "rope"]:69            raise ValueError(70                f"positional_bias_type must be 'alibi' or 'rope', got {positional_bias_type}"71            )72        self.static_residual = not static_residual73        self.no_moe_layers = no_moe_layers74        self.moe_bias = moe_bias75        self.balancing_coef = balancing_coef76        self.noisy_experts = noisy_experts77        self.high_rotations = high_rotations78        self.low_rotations = low_rotations79        self.positional_bias_type = positional_bias_type80        self.vocab_size = vocab_size81        self.hidden_size = hidden_size82        self.mlp_bias = mlp_bias83        self.num_hidden_layers = num_hidden_layers84        self.num_attention_heads = num_attention_heads85        self.num_kv_heads = num_kv_heads86        self.attention_dropout = attention_dropout87        self.rms_norm_eps = rms_norm_eps88        self.max_seq_len = max_seq_len89        self.use_cache = use_cache90        self.initializer_range = initializer_range91        self.embedding_dropout = embedding_dropout92        self.rms_affine = rms_affine93        self.output_hidden_states = output_hidden_states94        self.output_attentions = output_attentions95        self.layer_dropout = layer_dropout96        self.use_moe = use_moe97        self.moe_period = moe_period98        self.expert_size = expert_size99        self.shared_experts = shared_experts100        self.routed_experts = routed_experts101        self.token_experts = token_experts102        self.intermediate_size = intermediate_size103        self.attention_bias = attention_bias104        self.rope_base = rope_base105        self.head_size = head_size if head_size else hidden_size // num_attention_heads106        self.original_seq_len = (107            original_seq_len if original_seq_len is not None else max_seq_len108        )109 110        self.sliding_window_attention = sliding_window_attention111        self.sliding_window_context = sliding_window_context112        self.sliding_window_period = sliding_window_period113        if sliding_window_attention and sliding_window_context > max_seq_len:114            logger.warning(115                f"sliding_window_context more than max_seq_len, \116                    set sliding_window_context to {max_seq_len}"117            )118            self.sliding_window_context = max_seq_len119        if not sliding_window_attention:120            self.sliding_window_context = max_seq_len121 122        if self.head_size % 2 != 0 and self.positional_bias_type == "rope":123            raise ValueError("Head size should divided by 2")124 125        super().__init__(126            bos_token_id=bos_token_id,127            eos_token_id=eos_token_id,128            pad_token_id=pad_token_id,129            tie_word_embeddings=tie_word_embeddings,130            **kwargs,131        )132 133 134__all__ = ["SmalLmConfig"]135