| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from transformers.configuration_utils import PreTrainedConfig |
| 5 | |
| 6 | |
| 7 | class AfmoeConfig(PreTrainedConfig): |
| 8 | model_type = "afmoe" |
| 9 | |
| 10 | def __init__( |
| 11 | self, |
| 12 | vocab_size: int = 200_192, |
| 13 | hidden_size: int = 2048, |
| 14 | intermediate_size: int = 6144, |
| 15 | moe_intermediate_size: int = 1408, |
| 16 | num_hidden_layers: int = 32, |
| 17 | num_dense_layers: int = 1, |
| 18 | num_attention_heads: int = 16, |
| 19 | num_key_value_heads: int | None = None, |
| 20 | head_dim: int = 128, |
| 21 | hidden_act: str = "silu", |
| 22 | max_position_embeddings: int = 131072, |
| 23 | initializer_range: float = 0.02, |
| 24 | rms_norm_eps: float = 1e-5, |
| 25 | use_cache: bool = True, |
| 26 | tie_word_embeddings: bool = False, |
| 27 | rope_parameters: dict | None = None, |
| 28 | rope_scaling: dict | None = None, |
| 29 | num_experts: int = 64, |
| 30 | num_experts_per_tok: int = 6, |
| 31 | num_shared_experts: int = 2, |
| 32 | num_expert_groups: int = 1, |
| 33 | num_limited_groups: int = 1, |
| 34 | score_func: str = "sigmoid", |
| 35 | route_norm: bool = True, |
| 36 | route_scale: float = 1.0, |
| 37 | global_attn_every_n_layers: int = 4, |
| 38 | sliding_window: int = 2048, |
| 39 | layer_types: list[str] | None = None, |
| 40 | attention_dropout: float = 0.0, |
| 41 | mup_enabled: bool = False, |
| 42 | n_group: int = 1, |
| 43 | topk_group: int = 1, |
| 44 | **kwargs, |
| 45 | ): |
| 46 | self.vocab_size = vocab_size |
| 47 | self.hidden_size = hidden_size |
| 48 | self.intermediate_size = intermediate_size |
| 49 | self.num_hidden_layers = num_hidden_layers |
| 50 | self.num_dense_layers = num_dense_layers |
| 51 | self.num_attention_heads = num_attention_heads |
| 52 | self.num_key_value_heads = num_key_value_heads or num_attention_heads |
| 53 | self.head_dim = head_dim |
| 54 | self.hidden_act = hidden_act |
| 55 | self.max_position_embeddings = max_position_embeddings |
| 56 | self.initializer_range = initializer_range |
| 57 | self.rms_norm_eps = rms_norm_eps |
| 58 | self.use_cache = use_cache |
| 59 | rope_theta = kwargs.pop("rope_theta", 10000.0) |
| 60 | if rope_parameters is None: |
| 61 | rope_parameters = {"rope_type": "default", "rope_theta": rope_theta} |
| 62 | self.rope_parameters = rope_parameters |
| 63 | self.rope_scaling = rope_scaling |
| 64 | |
| 65 | self.moe_intermediate_size = moe_intermediate_size |
| 66 | self.num_experts = num_experts |
| 67 | self.num_experts_per_tok = num_experts_per_tok |
| 68 | self.num_shared_experts = num_shared_experts |
| 69 | self.num_expert_groups = num_expert_groups |
| 70 | self.num_limited_groups = num_limited_groups |
| 71 | self.score_func = score_func |
| 72 | self.route_norm = route_norm |
| 73 | self.route_scale = route_scale |
| 74 | |
| 75 | self.global_attn_every_n_layers = global_attn_every_n_layers |
| 76 | self.sliding_window = sliding_window |
| 77 | self.layer_types = layer_types |
| 78 | self.attention_dropout = attention_dropout |
| 79 | |
| 80 | self.mup_enabled = mup_enabled |
| 81 | self.n_group = n_group |
| 82 | self.topk_group = topk_group |
| 83 | |
| 84 | super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) |
| 85 | |
| 86 | |
| 87 | __all__ = ["AfmoeConfig"] |