| | 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) |
| |
|
