v-stone/testtesttest1632
012
1# coding=utf-8
2#
3# Copyright 2026 Xiaomi Corporation.
4# Copyright 2026 The HuggingFace Inc. team.
5#
6# Licensed under the Apache License, Version 2.0 (the "License");
7# you may not use this file except in compliance with the License.
8# You may obtain a copy of the License at
9#
10# http://www.apache.org/licenses/LICENSE-2.0
11#
12# Unless required by applicable law or agreed to in writing, software
13# distributed under the License is distributed on an "AS IS" BASIS,
14# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15# See the License for the specific language governing permissions and
16# limitations under the License.
17
18from transformers.configuration_utils import PretrainedConfig
19from transformers.modeling_rope_utils import rope_config_validation
20from transformers.utils import logging
21
22
23logger = logging.get_logger(__name__)
24
25
26_MIMOV2_ATTENTION_PROJECTION_LAYOUTS = {"split", "fused_qkv"}
27
28_MIMOV2_SPLIT_TP_PLAN = {
29 "layers.*.self_attn.q_proj": "colwise",
30 "layers.*.self_attn.k_proj": "colwise",
31 "layers.*.self_attn.v_proj": "colwise",
32 "layers.*.self_attn.o_proj": "rowwise",
33 "layers.*.mlp.gate_proj": "colwise",
34 "layers.*.mlp.up_proj": "colwise",
35 "layers.*.mlp.down_proj": "rowwise",
36}
37
38_MIMOV2_FUSED_QKV_TP_PLAN = {
39 "layers.*.self_attn.qkv_proj": "colwise",
40 "layers.*.self_attn.o_proj": "rowwise",
41 "layers.*.mlp.gate_proj": "colwise",
42 "layers.*.mlp.up_proj": "colwise",
43 "layers.*.mlp.down_proj": "rowwise",
44}
45
46_MIMOV2_PP_PLAN = {
47 "embed_tokens": (["input_ids"], ["inputs_embeds"]),
48 "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
49 "norm": (["hidden_states"], ["hidden_states"]),
50}
51
52
53class MiMoV2Config(PretrainedConfig):
54
55 model_type = "mimo_v2"
56 keys_to_ignore_at_inference = ["past_key_values"]
57
58 base_model_tp_plan = _MIMOV2_SPLIT_TP_PLAN
59 base_model_pp_plan = _MIMOV2_PP_PLAN
60
61 attribute_map = {
62 "num_local_experts": "n_routed_experts",
63 }
64
65 def __init__(
66 self,
67 vocab_size=151936,
68 hidden_size=4096,
69 intermediate_size=22016,
70 num_hidden_layers=32,
71 num_attention_heads=32,
72 num_key_value_heads=32,
73 hidden_act="silu",
74 max_position_embeddings=32768,
75 initializer_range=0.02,
76 layernorm_epsilon=1e-6,
77 use_cache=True,
78 tie_word_embeddings=False,
79 rope_theta=10000.0,
80 rope_scaling=None,
81 attention_dropout=0.0,
82 attention_bias=False,
83 attention_value_scale=None,
84 head_dim=None,
85 v_head_dim=None,
86 swa_num_attention_heads=None,
87 swa_num_key_value_heads=None,
88 swa_head_dim=None,
89 swa_v_head_dim=None,
90 swa_rope_theta=None,
91 sliding_window=None,
92 sliding_window_size=None,
93 add_full_attention_sink_bias=False,
94 add_swa_attention_sink_bias=False,
95 hybrid_block_size=None,
96 hybrid_layer_pattern=None,
97 partial_rotary_factor=1.0,
98 n_routed_experts=None,
99 moe_intermediate_size=None,
100 num_experts_per_tok=None,
101 routed_scaling_factor=None,
102 scoring_func="sigmoid",
103 topk_method="noaux_tc",
104 n_group=None,
105 topk_group=None,
106 norm_topk_prob=True,
107 moe_layer_freq=None,
108 attention_projection_layout="split",
109 **kwargs,
110 ):
111 rope_parameters = kwargs.pop("rope_parameters", None)
112 if rope_scaling is None and rope_parameters is not None:
113 rope_scaling = rope_parameters
114
115 if attention_projection_layout is None:
116 attention_projection_layout = "split"
117 if attention_projection_layout not in _MIMOV2_ATTENTION_PROJECTION_LAYOUTS:
118 raise ValueError(f"Unsupported MiMoV2 attention projection layout: {attention_projection_layout}")
119
120 self.attention_projection_layout = attention_projection_layout
121 self.base_model_tp_plan = (
122 _MIMOV2_FUSED_QKV_TP_PLAN.copy()
123 if attention_projection_layout == "fused_qkv"
124 else _MIMOV2_SPLIT_TP_PLAN.copy()
125 )
126 self.base_model_pp_plan = _MIMOV2_PP_PLAN.copy()
127
128 self.vocab_size = vocab_size
129 self.max_position_embeddings = max_position_embeddings
130 self.hidden_size = hidden_size
131 self.intermediate_size = intermediate_size
132 self.num_hidden_layers = num_hidden_layers
133 self.num_attention_heads = num_attention_heads
134
135 if num_key_value_heads is None:
136 num_key_value_heads = num_attention_heads
137 if num_attention_heads % num_key_value_heads != 0:
138 raise ValueError("num_attention_heads must be divisible by num_key_value_heads")
139
140 self.num_key_value_heads = num_key_value_heads
141 self.hidden_act = hidden_act
142 self.initializer_range = initializer_range
143 self.layernorm_epsilon = layernorm_epsilon
144 self.use_cache = use_cache
145 self.rope_theta = rope_theta
146 self.rope_scaling = rope_scaling
147 self.attention_dropout = attention_dropout
148 self.attention_bias = attention_bias
149 self.attention_value_scale = attention_value_scale
150
151 self.head_dim = head_dim if head_dim is not None else hidden_size // num_attention_heads
152 self.v_head_dim = v_head_dim if v_head_dim is not None else self.head_dim
153 self.swa_num_attention_heads = (
154 swa_num_attention_heads if swa_num_attention_heads is not None else num_attention_heads
155 )
156 self.swa_num_key_value_heads = (
157 swa_num_key_value_heads if swa_num_key_value_heads is not None else num_key_value_heads
158 )
159 if self.swa_num_attention_heads % self.swa_num_key_value_heads != 0:
160 raise ValueError("swa_num_attention_heads must be divisible by swa_num_key_value_heads")
161 self.swa_head_dim = swa_head_dim if swa_head_dim is not None else self.head_dim
162 self.swa_v_head_dim = swa_v_head_dim if swa_v_head_dim is not None else self.swa_head_dim
163 self.swa_rope_theta = swa_rope_theta if swa_rope_theta is not None else rope_theta
164
165 if sliding_window is None:
166 sliding_window = sliding_window_size
167 self.sliding_window = sliding_window
168 self.sliding_window_size = sliding_window_size if sliding_window_size is not None else sliding_window
169 self.add_full_attention_sink_bias = add_full_attention_sink_bias
170 self.add_swa_attention_sink_bias = add_swa_attention_sink_bias
171
172 if hybrid_block_size is not None and hybrid_layer_pattern is None:
173 hybrid_layer_pattern = [0 if ((i + 1) % hybrid_block_size == 0) else 1 for i in range(num_hidden_layers)]
174 elif hybrid_layer_pattern is None:
175 hybrid_layer_pattern = [0] * num_hidden_layers
176 if len(hybrid_layer_pattern) != num_hidden_layers:
177 raise ValueError("hybrid_layer_pattern length must match num_hidden_layers")
178 self.hybrid_block_size = hybrid_block_size
179 self.hybrid_layer_pattern = hybrid_layer_pattern
180
181 self.partial_rotary_factor = partial_rotary_factor
182
183 self.n_routed_experts = n_routed_experts
184 self.moe_intermediate_size = moe_intermediate_size if moe_intermediate_size is not None else intermediate_size
185 self.num_experts_per_tok = num_experts_per_tok
186 self.routed_scaling_factor = routed_scaling_factor
187 self.scoring_func = scoring_func
188 self.topk_method = topk_method
189 self.n_group = n_group
190 self.topk_group = topk_group
191 self.norm_topk_prob = norm_topk_prob
192 if isinstance(moe_layer_freq, int):
193 moe_layer_freq = [moe_layer_freq > 0 and i % moe_layer_freq == 0 for i in range(num_hidden_layers)]
194 elif moe_layer_freq is None:
195 moe_layer_freq = [False] * num_hidden_layers
196 if len(moe_layer_freq) != num_hidden_layers:
197 raise ValueError("moe_layer_freq length must match num_hidden_layers")
198 self.moe_layer_freq = moe_layer_freq
199
200 if self.rope_scaling is not None and "type" in self.rope_scaling:
201 self.rope_scaling["rope_type"] = self.rope_scaling["type"]
202 rope_config_validation(self)
203
204 super().__init__(
205 tie_word_embeddings=tie_word_embeddings,
206 **kwargs,
207 )
208
209__all__ = ["MiMoV2Config"]
210