1# Copyright 2025-2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""parser child class"""
16from hyper_parallel.auto_parallel.sapp_nd.nd.common.config import Config
17from hyper_parallel.auto_parallel.sapp_nd.nd.common.framework_parsers._cost_model_parser import _CostModelParser
18from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.size import Memory
19from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.logger import logger
20
21
22class CostModelParserMindformers(_CostModelParser):
23 """parser class for MindFormers format"""
24
25 def parse(self):
26 self.__config_parse_yaml()
27
28 def __config_parse_yaml_parallelism(self):
29 """MindFormer format for strategy"""
30 self.ccfg.has_op = self.config.parallel.enable_parallel_optimizer
31 op_cfg = self.config.parallel.parallel_optimizer_config
32 if op_cfg:
33 self.ccfg.has_grad_shard = op_cfg.gradient_accumulation_shard
34 self.ccfg.vocab_emb_dp = self.config.parallel_config.vocab_emb_dp
35 self.ccfg.tie_emb_out = self.config.model.model_config.tie_word_embeddings
36 self.ccfg.cp_algo = (
37 self.config.parallel_config.context_parallel_algo
38 if self.config.parallel_config.context_parallel_algo
39 else "colossalai_cp"
40 )
41 self.ccfg.d = max(
42 1, self.config.parallel_config.data_parallel
43 ) # Data parallel
44 self.ccfg.t = max(
45 1, self.config.parallel_config.model_parallel
46 ) # Tensor parallel
47 self.ccfg.p = max(
48 1, self.config.parallel_config.pipeline_stage
49 ) # Pipeline parallel
50 self.ccfg.cp = max(
51 1, self.config.parallel_config.context_parallel
52 ) # Context parallel
53 self.ccfg.ep = max(
54 1, self.config.parallel_config.expert_parallel
55 ) # Expert parallel
56 self.ccfg.sp = (
57 self.ccfg.t if self.config.parallel_config.use_seq_parallel else 1
58 ) # Sequence parallel factor
59 if self.ccfg.cp > 1 and self.ccfg.sp > 1:
60 logger.warning(
61 "sequence parallelism and context parallelism are both enabled"
62 )
63
64 # Interleaving
65 self.ccfg.n_s_split = 1 # seqpipe
66 self.ccfg.pp_sched = "1f1b"
67 self.ccfg.vp = max(1, self.config.model.model_config.pp_interleave_num)
68 if self.config.parallel.pipeline_config:
69 if self.config.parallel.pipeline_config.pipeline_interleave:
70 self.ccfg.n_s_split = max(
71 1, self.config.parallel_config.seq_split_num
72 )
73 if self.config.parallel.pipeline_config.pipeline_scheduler:
74 self.ccfg.pp_sched = (
75 self.config.parallel.pipeline_config.pipeline_scheduler
76 )
77
78 def __config_parse_yaml_hyperparameters(self):
79 """MindFormer format for hyperparams"""
80 self.ccfg.n_lay = (
81 self.config.model.model_config.num_layers
82 if self.config.model.model_config.num_layers
83 else self.config.model.model_config.num_hidden_layers
84 )
85 if self.config.model.model_config.is_encoder_decoder:
86 self.ccfg.n_lay *= 2
87 self.ccfg.h = self.config.model.model_config.hidden_size # Hidden size
88 self.ccfg.hff = (
89 self.config.model.model_config.intermediate_size
90 if self.config.model.model_config.intermediate_size
91 else self.init_hff()
92 ) # Expanded hidden size
93 self.ccfg.v = self.config.model.model_config.vocab_size # Vocabulary size
94 self.ccfg.s = self.config.model.model_config.seq_length # Sequence length
95 self.ccfg.a = (
96 self.config.model.model_config.num_heads
97 ) # Number of attention (query) heads
98 if not self.ccfg.a:
99 self.ccfg.a = self.config.model.model_config.num_attention_heads
100 self.ccfg.n_kv = self.config.model.model_config.n_kv_heads # Number of keys - values heads
101 if not self.ccfg.n_kv:
102 self.ccfg.n_kv = self.config.model.model_config.num_key_value_heads
103 if not self.ccfg.n_kv:
104 self.ccfg.n_kv = self.ccfg.a
105 self.ccfg.dh = self.ccfg.h / self.ccfg.a # Per head dimension
106 self.ccfg.dc_kv = (
107 self.config.model.model_config.kv_lora_rank
108 if self.config.model.model_config.kv_lora_rank
109 else 0
110 ) # KV compression dimension
111 self.ccfg.dc_q = (
112 self.config.model.model_config.q_lora_rank
113 if self.config.model.model_config.q_lora_rank
114 else 0
115 ) # Q compression dimension
116 self.ccfg.dhr = (
117 self.config.model.model_config.qk_rope_head_dim
118 if self.config.model.model_config.qk_rope_head_dim
119 else 0
120 ) # decoupled QK per head dimension
121
122 # Microbatch infos
123 self.ccfg.b = max(
124 1, self.config.runner_config.batch_size
125 ) # Microbatch size
126 self.ccfg.m = (
127 self.config.parallel_config.micro_batch_num
128 ) # Number of microbatches
129 if self.ccfg.m <= 0:
130 logger.warning("num_micro is negative")
131
132 def __config_parse_yaml_moe(self):
133 """MindFormer format for MoE infos"""
134 self.ccfg.n_exp, self.ccfg.n_chosen_exp, self.ccfg.n_shared_exp = 1, 1, 0
135 self.ccfg.hff_exp, self.ccfg.cap_fact = self.ccfg.hff, 1
136 self.ccfg.t_exp, self.ccfg.d_exp = self.ccfg.t, self.ccfg.d
137 if self.config.moe_config:
138 self.ccfg.n_exp = max(
139 1, self.config.moe_config.expert_num
140 ) # Total number of experts
141 self.ccfg.n_chosen_exp = max(
142 1, self.config.moe_config.num_experts_chosen
143 ) # Number of chosen experts
144 self.ccfg.n_shared_exp = self.config.moe_config.shared_expert_num
145 if self.config.moe_config.moe_intermediate_size:
146 self.ccfg.hff_exp = self.config.moe_config.moe_intermediate_size
147 self.ccfg.cap_fact = max(
148 1, self.config.moe_config.capacity_factor
149 ) # Capacity factor
150 self.ccfg.k_1st_dense = self.config.moe_config.first_k_dense_replace
151 self.ccfg.etp = self.config.moe_config.expert_model_parallel
152 self.config_dp_tp_exp(self.ccfg)
153 self.ccfg.gmm = self.config.moe_config.use_gmm
154 else:
155 cfg = self.config.model.model_config
156 self.ccfg.n_exp = max(self.ccfg.n_exp, cfg.n_routed_experts)
157 self.ccfg.n_chosen_exp = max(self.ccfg.n_chosen_exp, cfg.num_experts_per_tok)
158 self.ccfg.n_shared_exp = max(self.ccfg.n_shared_exp, cfg.n_shared_experts)
159 if cfg.moe_intermediate_size:
160 self.ccfg.hff_exp = cfg.moe_intermediate_size
161 self.ccfg.k_1st_dense = max(self.ccfg.k_1st_dense, cfg.first_k_dense_replace)
162 self.config_dp_tp_exp(self.ccfg)
163 self.ccfg.gmm = cfg.moe_grouped_gemm
164
165 def __config_parse_yaml_op_recompute(self):
166 """MindFormer format for select recompute"""
167 # [HYPOTHESIS]
168 self.ccfg.rec_op = Config(
169 {}
170 ) # recomputed operators (selective recompute only)
171 self.ccfg.rec_op.attBMM = int(
172 not (
173 self.config.recompute_config.select_recompute
174 and not self.ccfg.has_fa
175 and self.ccfg.sp > 1
176 )
177 )
178 self.ccfg.rec_op.headCast = int(
179 not (self.config.recompute_config.select_recompute and self.ccfg.has_fa)
180 )
181 self.ccfg.rec_op.dropout = 1
182 self.ccfg.rec_op.softmax = int(
183 not (
184 self.config.recompute_config.select_recompute
185 and not self.ccfg.has_fa
186 )
187 )
188 self.ccfg.rec_op.normOp = int(
189 not (self.config.recompute_config.select_recompute and self.ccfg.sp > 1)
190 )
191 self.ccfg.rec_op.gather = int(
192 not (
193 self.config.recompute_config.select_comm_recompute
194 and self.ccfg.sp > 1
195 )
196 )
197 self.ccfg.rec_op.ffAct = int(
198 not (self.config.recompute_config.select_recompute and self.ccfg.sp > 1)
199 )
200
201 def config_shard_emb(self):
202 """Configure embedding and output activation sharding."""
203 self.ccfg.shard_embed = (
204 self.ccfg.d
205 if (self.ccfg.vocab_emb_dp and self.ccfg.p == 1)
206 else (self.ccfg.t * self.ccfg.d)
207 )
208
209 def __config_parse_yaml(self):
210 """MindFormer format for unimodal"""
211 self.ccfg.config_format = "yaml"
212 self.ccfg.model_name = self.config.trainer.model_name
213 self.ccfg.device_capacity = Memory.from_string(
214 self.config.context.max_device_memory
215 )
216 # (
217 # float(self.config.context.max_device_memory[:-2])
218 # * 1024
219 # * 1024
220 # * 1024
221 # )
222 self.ccfg.has_fa = self.config.model.model_config.use_flash_attention
223 # self.ccfg.vp_less_mem = False
224 self.ccfg.has_clip = self.config.runner_wrapper.use_clip_grad
225 self.ccfg.gmm = False
226 self.ccfg.freeze = False
227 op_cfg = self.config.parallel.parallel_optimizer_config
228 if op_cfg:
229 self.ccfg.op_weight_shard = op_cfg.optimizer_weight_shard_size
230 self.ccfg.optimizer = self.config.optimizer.type
231 self.ccfg.multiple_of = (
232 self.config.model.model_config.multiple_of
233 if self.config.model.model_config.multiple_of
234 else 1
235 )
236 self.ccfg.fdm = (
237 self.config.model.model_config.ffn_dim_multiplier
238 if self.config.model.model_config.ffn_dim_multiplier
239 else 1
240 )
241 self.__config_parse_yaml_parallelism()
242 self.__config_parse_yaml_hyperparameters()
243 self.__config_parse_yaml_moe()
244
245 # FP byte storages
246 self.ccfg.bytes_p = self.ccfg.fp_bytes(
247 self.config.model.model_config.param_init_type,
248 ) # parameters
249 if not self.ccfg.bytes_p:
250 self.ccfg.bytes_p = self.ccfg.fp_bytes(
251 self.config.model.model_config.params_dtype
252 )
253 self.ccfg.bytes_compute = self.ccfg.fp_bytes(
254 self.config.model.model_config.compute_dtype
255 ) # activations
256 self.ccfg.bytes_softmax = self.ccfg.fp_bytes(
257 self.config.model.model_config.softmax_compute_type
258 ) # softmax output
259 if not self.ccfg.bytes_softmax:
260 self.ccfg.bytes_softmax = self.ccfg.fp_bytes(
261 self.config.model.model_config.softmax_compute_dtype
262 )
263 if not self.ccfg.bytes_p:
264 raise AttributeError("bytes_p not positive")
265 if not self.ccfg.bytes_compute:
266 raise AttributeError("bytes_compute not positive")
267
268 # Optimizer parallel factors
269 if self.ccfg.op_weight_shard:
270 self.ccfg.os_max_shard = self.ccfg.op_weight_shard
271 elif self.ccfg.has_op:
272 self.ccfg.os_max_shard = self.ccfg.d * self.ccfg.t
273 else:
274 self.ccfg.os_max_shard = 1
275 self.config_optimizer_shard(self.ccfg)
276
277 # Other factors
278 self.config_shard_emb()
279 self.ccfg.shard_output_activ = 1
280 self.ccfg.shard_recompute_input = (
281 self.ccfg.t
282 if self.config.recompute_config.recompute_slice_activation
283 else 1
284 )
285 self.ccfg.s_fa = (
286 self.ccfg.s
287 if not self.config.model.model_config.use_flash_attention
288 else self.ccfg.s / self.ccfg.a
289 ) # flash attention factor [HYPOTHESIS]
290 self.config_comm_flag(self.ccfg)
291 self.ccfg.gbs = self.ccfg.b * self.ccfg.d * self.ccfg.m
292 self.ccfg.n_mtp = (
293 self.config.model.model_config.mtp_depth
294 if self.config.model.model_config.mtp_depth
295 else 0
296 )
297 if not self.ccfg.n_mtp:
298 self.ccfg.n_mtp = (
299 self.config.model.model_config.num_nextn_predict_layers
300 )
301 self.ccfg.is_mtp_in_offset = False
302
303 # Layer custom config
304 # [(num layers selected, layer custom config to apply)]
305 self.ccfg.layer_custom_config = [(self.ccfg.n_lay + self.ccfg.n_mtp, None)]
306
307 self.__config_parse_yaml_op_recompute()
308 # By default, 100% of layers use a unique custom config (if specified)
309 self.ccfg.offset = self.config.model.model_config.offset
310 self.ccfg.sel_rec = self.config.recompute_config.select_recompute
311 self.ccfg.full_rec = self.config.recompute_config.recompute
312 self.ccfg.overwrite_eval_functions = {}