| from transformers import PretrainedConfig |
| from typing import List |
|
|
|
|
| class LMConfig(PretrainedConfig): |
| model_type = "minimind" |
|
|
| def __init__( |
| self, |
| dim: int = 512, |
| n_layers: int = 8, |
| n_heads: int = 8, |
| n_kv_heads: int = 2, |
| vocab_size: int = 6400, |
| hidden_dim: int = None, |
| multiple_of: int = 64, |
| norm_eps: float = 1e-5, |
| max_seq_len: int = 8192, |
| rope_theta: int = 1e6, |
| dropout: float = 0.0, |
| flash_attn: bool = True, |
| |
| |
| |
| |
| use_moe: bool = False, |
| |
| num_experts_per_tok: int = 2, |
| n_routed_experts: int = 4, |
| n_shared_experts: bool = True, |
| scoring_func: str = 'softmax', |
| aux_loss_alpha: float = 0.1, |
| seq_aux: bool = True, |
| norm_topk_prob: bool = True, |
| **kwargs, |
| ): |
| self.dim = dim |
| self.n_layers = n_layers |
| self.n_heads = n_heads |
| self.n_kv_heads = n_kv_heads |
| self.vocab_size = vocab_size |
| self.hidden_dim = hidden_dim |
| self.multiple_of = multiple_of |
| self.norm_eps = norm_eps |
| self.max_seq_len = max_seq_len |
| self.rope_theta = rope_theta |
| self.dropout = dropout |
| self.flash_attn = flash_attn |
| |
| |
| |
| |
| self.use_moe = use_moe |
| self.num_experts_per_tok = num_experts_per_tok |
| self.n_routed_experts = n_routed_experts |
| self.n_shared_experts = n_shared_experts |
| self.scoring_func = scoring_func |
| self.aux_loss_alpha = aux_loss_alpha |
| self.seq_aux = seq_aux |
| self.norm_topk_prob = norm_topk_prob |
| super().__init__(**kwargs) |
|
|
