File size: 2,921 Bytes
70038b6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
from transformers.configuration_utils import PretrainedConfig


class DynamicMindMoEConfig(PretrainedConfig):
    """DynamicMind-MoE: sparse mixture-of-experts variant of DynamicMind-Mini.



    The dense MLP (intermediate 768) is replaced by one always-on shared expert

    plus `num_routed_experts` fine-grained experts (intermediate 256), of which

    `num_experts_per_token` are selected. Shared + top-2 reproduces the dense

    layer's exact active parameter count, so inference cost per token is

    unchanged while total capacity grows ~3.4x.

    """

    model_type = "dynamicmind_moe"

    def __init__(

        self,

        vocab_size=8192,

        hidden_size=256,

        intermediate_size=768,        # kept for dense layers / upcycling source

        moe_intermediate_size=256,    # per-expert width (768 / 3)

        num_hidden_layers=9,

        num_attention_heads=8,

        num_key_value_heads=2,

        num_routed_experts=14,

        num_shared_experts=1,

        num_experts_per_token=2,

        first_k_dense_layers=0,       # keep the first K blocks dense if desired

        norm_topk_prob=True,

        router_aux_loss_coef=0.01,

        router_z_loss_coef=0.001,

        router_bias_update_rate=0.001,  # aux-loss-free balancing (DeepSeek-V3)

        use_aux_loss_free_balancing=True,

        max_position_embeddings=1024,

        rms_norm_eps=1e-5,

        rope_theta=10000.0,

        attention_dropout=0.0,

        tie_word_embeddings=True,

        bos_token_id=0,

        eos_token_id=0,

        pad_token_id=1,

        **kwargs,

    ):
        super().__init__(
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            pad_token_id=pad_token_id,
            tie_word_embeddings=tie_word_embeddings,
            **kwargs,
        )

        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.moe_intermediate_size = moe_intermediate_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.num_key_value_heads = num_key_value_heads

        self.num_routed_experts = num_routed_experts
        self.num_shared_experts = num_shared_experts
        self.num_experts_per_token = num_experts_per_token
        self.first_k_dense_layers = first_k_dense_layers
        self.norm_topk_prob = norm_topk_prob

        self.router_aux_loss_coef = router_aux_loss_coef
        self.router_z_loss_coef = router_z_loss_coef
        self.router_bias_update_rate = router_bias_update_rate
        self.use_aux_loss_free_balancing = use_aux_loss_free_balancing

        self.max_position_embeddings = max_position_embeddings
        self.rms_norm_eps = rms_norm_eps
        self.rope_theta = rope_theta
        self.attention_dropout = attention_dropout