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.framework_parsers._cost_model_parser import _CostModelParser
17from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.size import Memory
18from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.logger import logger
19from hyper_parallel.auto_parallel.sapp_nd.nd.common.derive import derive
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.emb_dp_sharded = True
36 self.ccfg.tie_emb_out = self.config.model.model_config.tie_word_embeddings
37 self.ccfg.cp_algo = (
38 self.config.parallel_config.context_parallel_algo
39 if self.config.parallel_config.context_parallel_algo
40 else "colossalai_cp"
41 )
42 self.ccfg.d = max(
43 1, self.config.parallel_config.data_parallel
44 ) # Data parallel
45 self.ccfg.t = max(
46 1, self.config.parallel_config.model_parallel
47 ) # Tensor parallel
48 self.ccfg.p = max(
49 1, self.config.parallel_config.pipeline_stage
50 ) # Pipeline parallel
51 self.ccfg.cp = max(
52 1, self.config.parallel_config.context_parallel
53 ) # Context parallel
54 self.ccfg.ep = max(
55 1, self.config.parallel_config.expert_parallel
56 ) # Expert parallel
57 self.ccfg.sequence_parallel = bool(self.config.parallel_config.use_seq_parallel)
58 if self.ccfg.cp > 1 and self.ccfg.sequence_parallel and self.ccfg.t > 1:
59 logger.warning(
60 "sequence parallelism and context parallelism are both enabled"
61 )
62
63 # Interleaving
64 self.ccfg.n_s_split = 1 # seqpipe
65 self.ccfg.pp_sched = "1f1b"
66 self.ccfg.vp = max(1, self.config.model.model_config.pp_interleave_num)
67 if self.config.parallel.pipeline_config:
68 if self.config.parallel.pipeline_config.pipeline_interleave:
69 self.ccfg.n_s_split = max(
70 1, self.config.parallel_config.seq_split_num
71 )
72 if self.config.parallel.pipeline_config.pipeline_scheduler:
73 self.ccfg.pp_sched = (
74 self.config.parallel.pipeline_config.pipeline_scheduler
75 )
76
77 def __config_parse_yaml_hyperparameters(self):
78 """MindFormer format for hyperparams"""
79 self.ccfg.n_lay = (
80 self.config.model.model_config.num_layers
81 if self.config.model.model_config.num_layers
82 else self.config.model.model_config.num_hidden_layers
83 )
84 if self.config.model.model_config.is_encoder_decoder:
85 self.ccfg.n_lay *= 2
86 self.ccfg.h = self.config.model.model_config.hidden_size # Hidden size
87 self.ccfg.hff = (
88 self.config.model.model_config.intermediate_size
89 if self.config.model.model_config.intermediate_size
90 else self.init_hff()
91 ) # Expanded hidden size
92 self.ccfg.v = self.config.model.model_config.vocab_size # Vocabulary size
93 self.ccfg.s = self.config.model.model_config.seq_length # Sequence length
94 self.ccfg.a = (
95 self.config.model.model_config.num_heads
96 ) # Number of attention (query) heads
97 if not self.ccfg.a:
98 self.ccfg.a = self.config.model.model_config.num_attention_heads
99 self.ccfg.n_kv = self.config.model.model_config.n_kv_heads # Number of keys - values heads
100 if not self.ccfg.n_kv:
101 self.ccfg.n_kv = self.config.model.model_config.num_key_value_heads
102 if not self.ccfg.n_kv:
103 self.ccfg.n_kv = self.ccfg.a
104 self.ccfg.dh = self.ccfg.h / self.ccfg.a # Per head dimension
105 # An MLA model's value-head width, which its family prices dh at.
106 self.ccfg.v_head_dim = self.config.model.model_config.v_head_dim or None
107 # Its heads' non-rotary key width, the value heads' unless stated.
108 self.ccfg.qk_nope_head_dim = self.config.model.model_config.qk_nope_head_dim or None
109 self.ccfg.dc_kv = (
110 self.config.model.model_config.kv_lora_rank
111 if self.config.model.model_config.kv_lora_rank
112 else 0
113 ) # KV compression dimension
114 self.ccfg.dc_q = (
115 self.config.model.model_config.q_lora_rank
116 if self.config.model.model_config.q_lora_rank
117 else 0
118 ) # Q compression dimension
119 self.ccfg.dhr = (
120 self.config.model.model_config.qk_rope_head_dim
121 if self.config.model.model_config.qk_rope_head_dim
122 else 0
123 ) # decoupled QK per head dimension
124 # Whether each head's queries and keys are normalized, as MindFormers'
125 # Qwen3 states it.
126 self.ccfg.qk_norm = bool(self.config.model.model_config.qk_layernorm)
127
128 # Microbatch infos
129 self.ccfg.b = max(
130 1, self.config.runner_config.batch_size
131 ) # Microbatch size
132 self.ccfg.m = (
133 self.config.parallel_config.micro_batch_num
134 ) # Number of microbatches
135 if self.ccfg.m <= 0:
136 logger.warning("num_micro is negative")
137
138 def __config_parse_yaml_moe(self):
139 """MindFormer format for MoE infos"""
140 self.ccfg.n_exp, self.ccfg.n_chosen_exp, self.ccfg.n_shared_exp = 1, 1, 0
141 self.ccfg.hff_exp, self.ccfg.cap_fact = self.ccfg.hff, 1
142 self.ccfg.t_exp, self.ccfg.d_exp = self.ccfg.t, self.ccfg.d
143 if self.config.moe_config:
144 self.ccfg.n_exp = max(
145 1, self.config.moe_config.expert_num
146 ) # Total number of experts
147 self.ccfg.n_chosen_exp = max(
148 1, self.config.moe_config.num_experts_chosen
149 ) # Number of chosen experts
150 self.ccfg.n_shared_exp = self.config.moe_config.shared_expert_num
151 if self.config.moe_config.moe_intermediate_size:
152 self.ccfg.hff_exp = self.config.moe_config.moe_intermediate_size
153 self.ccfg.cap_fact = max(
154 1, self.config.moe_config.capacity_factor
155 ) # Capacity factor
156 self.ccfg.k_1st_dense = self.config.moe_config.first_k_dense_replace
157 self.ccfg.etp = self.config.moe_config.expert_model_parallel
158 self.ccfg.gmm = self.config.moe_config.use_gmm
159 else:
160 cfg = self.config.model.model_config
161 self.ccfg.n_exp = max(self.ccfg.n_exp, cfg.n_routed_experts)
162 self.ccfg.n_chosen_exp = max(self.ccfg.n_chosen_exp, cfg.num_experts_per_tok)
163 self.ccfg.n_shared_exp = max(self.ccfg.n_shared_exp, cfg.n_shared_experts)
164 if cfg.moe_intermediate_size:
165 self.ccfg.hff_exp = cfg.moe_intermediate_size
166 self.ccfg.k_1st_dense = max(self.ccfg.k_1st_dense, cfg.first_k_dense_replace)
167 self.ccfg.gmm = cfg.moe_grouped_gemm
168
169 def __config_parse_yaml_recompute(self):
170 """MindFormer format for recompute"""
171 rc = self.config.recompute_config
172 self.ccfg.sel_rec = rc.select_recompute
173 self.ccfg.sel_comm_rec = rc.select_comm_recompute
174 self.ccfg.sel_rec_rule = "mindformers"
175 self.ccfg.full_rec = rc.recompute
176 self.ccfg.recompute_slice_activation = bool(rc.recompute_slice_activation)
177
178 def __config_parse_yaml_fp_bytes(self):
179 """FP byte storages of parameters, activations and softmax outputs."""
180 self.ccfg.bytes_p = self.ccfg.fp_bytes(
181 self.config.model.model_config.param_init_type,
182 ) # parameters
183 if not self.ccfg.bytes_p:
184 self.ccfg.bytes_p = self.ccfg.fp_bytes(
185 self.config.model.model_config.params_dtype
186 )
187 self.ccfg.bytes_compute = self.ccfg.fp_bytes(
188 self.config.model.model_config.compute_dtype
189 ) # activations
190 self.ccfg.bytes_softmax = self.ccfg.fp_bytes(
191 self.config.model.model_config.softmax_compute_type
192 ) # softmax output
193 if not self.ccfg.bytes_softmax:
194 self.ccfg.bytes_softmax = self.ccfg.fp_bytes(
195 self.config.model.model_config.softmax_compute_dtype
196 )
197 if not self.ccfg.bytes_p:
198 raise AttributeError("bytes_p not positive")
199 if not self.ccfg.bytes_compute:
200 raise AttributeError("bytes_compute not positive")
201
202 def __config_parse_yaml_optimizer_shard(self):
203 """Optimizer parallel factors."""
204 if self.ccfg.op_weight_shard:
205 self.ccfg.os_max_shard = self.ccfg.op_weight_shard
206 elif self.ccfg.has_op:
207 self.ccfg.os_max_shard = self.ccfg.d
208 else:
209 self.ccfg.os_max_shard = 1
210
211 def __config_parse_yaml(self):
212 """MindFormer format for unimodal"""
213 self.ccfg.config_format = "yaml"
214 self.ccfg.model_name = self.config.trainer.model_name
215 self.config_op_counts(self.ccfg)
216 self.ccfg.device_capacity = Memory.from_string(
217 self.config.context.max_device_memory
218 )
219 # (
220 # float(self.config.context.max_device_memory[:-2])
221 # * 1024
222 # * 1024
223 # * 1024
224 # )
225 self.ccfg.has_fa = self.config.model.model_config.use_flash_attention
226 # self.ccfg.vp_less_mem = False
227 self.ccfg.has_clip = self.config.runner_wrapper.use_clip_grad
228 self.ccfg.gmm = False
229 self.ccfg.freeze = False
230 op_cfg = self.config.parallel.parallel_optimizer_config
231 if op_cfg:
232 self.ccfg.op_weight_shard = op_cfg.optimizer_weight_shard_size
233 # Read from the section's own fields: a missing section reads as 0,
234 # whose .type raised, where the HyperParallel parser takes AdamW.
235 optimizer = vars(self.config).get("optimizer")
236 self.state_optimizer(self.ccfg, getattr(optimizer, "__dict__", {}).get("type"))
237 self.ccfg.multiple_of = (
238 self.config.model.model_config.multiple_of
239 if self.config.model.model_config.multiple_of
240 else 1
241 )
242 self.ccfg.fdm = (
243 self.config.model.model_config.ffn_dim_multiplier
244 if self.config.model.model_config.ffn_dim_multiplier
245 else 1
246 )
247 self.__config_parse_yaml_parallelism()
248 self.__config_parse_yaml_hyperparameters()
249 self.__config_parse_yaml_moe()
250 self.__config_parse_yaml_fp_bytes()
251 self.__config_parse_yaml_optimizer_shard()
252 self.ccfg.gbs = self.ccfg.b * self.ccfg.d * self.ccfg.m
253 self.ccfg.n_mtp = (
254 self.config.model.model_config.mtp_depth
255 if self.config.model.model_config.mtp_depth
256 else 0
257 )
258 if not self.ccfg.n_mtp:
259 self.ccfg.n_mtp = (
260 self.config.model.model_config.num_nextn_predict_layers
261 )
262 self.ccfg.is_mtp_in_offset = False
263
264 self.config_layer_stack(self.ccfg)
265
266 # By default, 100% of layers use a unique custom config (if specified)
267 self.ccfg.offset = self.config.model.model_config.offset
268 self.__config_parse_yaml_recompute()
269 self.ccfg.overwrite_eval_functions = {}
270 derive(self.ccfg)