« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +0800
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, YamlObject
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
22class CostModelParserMindspeed(_CostModelParser):
23 """parser class for MindSpeed format"""
25 def parse(self):
26 self.__config_parse_json_multimodals()
28 # MindSpeed (multimodal)
29 def __config_parse_json_multimodals(self):
30 """MindSpeed format for multimodal"""
31 self.ccfg.device_capacity = Memory.from_gb(55) # 55 * 1024 * 1024 * 1024
32 self.ccfg.model_name = self.config.model_id
33 # Assume it exists a hook module with the same name as model_name
34 self.ccfg.n_lay = 0 # SUM ALL
35 self.ccfg.pp_sched = "1f1b"
36 self.ccfg.p = self.config.tmp.pp # PP
37 self.ccfg.m = self.ccfg.p
38 self.ccfg.b = self.config.tmp.mbs
39 self.ccfg.d = self.config.tmp.dp # DP
40 self.ccfg.t = self.config.tmp.tp # TP
41 self.ccfg.os_max_shard = self.ccfg.d * self.ccfg.t
42 self.ccfg.cp = self.config.tmp.cp # CP
43 self.ccfg.vp = self.config.tmp.vpp # VPP
44 self.ccfg.ep = self.config.tmp.ep # EP
45 ccfgs = self.__search_and_parse_mods_ccfg(self.ccfg.config)
46 self.ccfg.multimodal = len(ccfgs) > 1
47 if self.ccfg.multimodal:
48 if not self.ccfg.hooks_dict:
49 raise TypeError(
50 "Currently for multimodals, 'hook_cls' required for evaluator"
51 )
52 missing_hooks = set(ccfgs.keys()) - set(self.ccfg.hooks_dict.keys())
53 if missing_hooks:
54 raise TypeError(
55 f"Missing hooks for submodules {list(missing_hooks)}"
56 )
57 self.ccfg.mm_ccfgs = ccfgs
58 self.ccfg.mm_order = list(ccfgs.keys())
60 # Update each mod's offset, layer_custom_config, pp_partition
61 for m in self.ccfg.mm_ccfgs:
62 cc = self.ccfg.mm_ccfgs[m]
63 num_layer_per_stage = max(1, cc.n_lay // self.ccfg.p // self.ccfg.vp)
64 if cc.pp_partition:
65 # print("here",m,cc.pp_partition,num_layer_per_stage)
66 # Making sure the format is offset[vp][p] like in MF
67 if isinstance(cc.pp_partition[0], list):
68 cc.offset = [
69 [
70 cc.pp_partition[v_idx][idx] - num_layer_per_stage
71 for idx in range(self.ccfg.p)
72 ]
73 for v_idx in range(self.ccfg.vp)
74 ]
75 else:
76 cc.offset = [
77 [
78 cc.pp_partition[idx] // self.ccfg.vp
79 - num_layer_per_stage
80 for idx in range(self.ccfg.p)
81 ]
82 for v_idx in range(self.ccfg.vp)
83 ]
84 else:
85 self.__complete_unimodal_pp_plan(m, cc, num_layer_per_stage)
86 self.ccfg.overwrite_eval_functions = {}
88 def __complete_unimodal_pp_plan(self, m, cc, num_layer_per_stage):
89 """Try to follow previous pp plan"""
90 stage_insert_idx, chunk_insert_idx = 0, 0
91 cc.offset = [
92 [-num_layer_per_stage for _ in range(self.ccfg.p)]
93 for _ in range(self.ccfg.vp)
94 ]
95 previous_mod_idx = (
96 self.ccfg.mm_order.index(m) - 1
97 ) # look for previous mod pp partition
98 if previous_mod_idx >= 0:
99 previous_mod_partition = self.ccfg.mm_ccfgs[
100 self.ccfg.mm_order[previous_mod_idx]
101 ].pp_partition
102 if isinstance(previous_mod_partition[0], list): # vpp
103 put = False
104 for v_idx in range(self.ccfg.vp - 1, -1, -1):
105 for s_idx in range(self.ccfg.p - 1, -1, -1):
106 if previous_mod_partition[v_idx][s_idx]:
107 stage_insert_idx = s_idx
108 chunk_insert_idx = v_idx
109 put = True
110 break
111 if put:
112 break
113 else:
114 stage_insert_idx = self.ccfg.p - 1
115 for p in previous_mod_partition[::-1]:
116 if p > 0:
117 break
118 stage_insert_idx = max(0, stage_insert_idx - 1)
119 cc.offset[chunk_insert_idx][stage_insert_idx] = (
120 cc.n_lay - num_layer_per_stage
121 )
122 cc.pp_partition = [
123 [
124 num_layer_per_stage + cc.offset[v_idx][idx]
125 for idx in range(self.ccfg.p)
126 ]
127 for v_idx in range(self.ccfg.vp)
128 ]
129 cc.p, cc.vp = self.ccfg.p, self.ccfg.vp
130 cc.full_rec = (
131 self.ccfg.mm_ccfgs[self.ccfg.mm_order[previous_mod_idx]].full_rec is True
132 )
133 cc.sel_rec = False
135 def __search_and_parse_mods_ccfg(self, field):
136 """extract multimodal submodules (MindSpeed format)"""
137 res = {}
138 for _, v in vars(field).items():
139 if isinstance(v, YamlObject):
140 res.update(self.__search_and_parse_mods_ccfg(v))
141 if v.model_id:
142 logger.info("Detected model config: %s", v.model_id)
143 res[v.model_id] = self.__config_parse_json(
144 v
145 ) # Build cost model variable foreach configs
146 return res
148 def __config_parse_json_parallelism(self, cc, mod):
149 """MindSpeed format for parallelism"""
150 cc.t = max(cc.tensor_model_parallel_size, self.config.tmp.tp)
151 cc.p = max(cc.pipeline_model_parallel_size, self.config.tmp.pp)
152 cc.cp = self.config.tmp.cp
153 cc.d = self.config.tmp.dp
154 cc.ep = max(cc.expert_model_parallel_size, self.config.tmp.ep)
155 cc.sp = cc.t if mod.sequence_parallel else 1
156 if cc.cp > 1 and cc.sp > 1:
157 logger.warning(
158 "sequence parallelism and context parallelism are both enabled"
159 )
160 cc.pp_partition = (
161 mod.pipeline_num_layers
162 ) # Offset regardless of even distribution
164 # Interleaving
165 cc.n_s_split = 1 # seqpipe
166 cc.pp_sched = "1f1b"
167 cc.vp = self.config.tmp.vpp # VPP
169 def __config_parse_json_hyperparameters(self, cc, mod):
170 """MindSpeed format for hyperparams"""
171 cc.n_lay = mod.num_layers
172 cc.h = mod.hidden_size
173 cc.hff = mod.ffn_hidden_size
174 cc.v = mod.vocab_size
175 cc.s = self.config.tmp.seqlen # SEQ_LEN
176 cc.a = mod.num_attention_heads
177 cc.n_kv = (
178 mod.num_query_groups if mod.num_query_groups else cc.a
179 ) # NOT SURE
180 cc.dh = (
181 mod.kv_channels if mod.kv_channels else (cc.h / cc.a)
182 ) # Per head dimension
183 # print(cc.a, cc.h, cc.dh)
184 cc.dc_kv = mod.k_lora_rank # KV compression dimension #NOT SURE
185 cc.dc_q = mod.q_lora_rank # Q compression dimension
186 cc.dhr = mod.qk_rope_head_dim # decoupled QK per head dimension
188 def __config_parse_json_moe(self, cc, mod):
189 """MindSpeed format for MoE infos"""
190 cc.n_exp = max(1, mod.num_moe_experts)
191 cc.n_chosen_exp = max(1, mod.moe_router_topk)
192 cc.n_shared_exp = mod.n_shared_exp
193 cc.cap_fact = 1
194 cc.t_exp, cc.d_exp = cc.t, cc.d
195 cc.hff_exp = (
196 mod.moe_intermediate_size if mod.moe_intermediate_size else cc.hff
197 )
198 cc.k_1st_dense = mod.first_k_dense_replace
199 # temporary
200 cc.etp = self.config.tmp.etp # ETP
202 def __config_parse_json_op_recompute(self, cc):
203 """MindSpeed format for select recompute"""
204 cc.rec_op = Config(
205 {}
206 ) # recomputed operators (selective recompute only)
207 cc.rec_op.attBMM = 1
208 cc.rec_op.headCast = 1
209 cc.rec_op.dropout = 1
210 cc.rec_op.softmax = 1
211 cc.rec_op.normOp = 1
212 cc.rec_op.gather = 1
213 cc.rec_op.ffAct = 1
215 def __config_parse_json(self, mod):
216 """MindSpeed format for unimodal"""
217 # def mod_hook(M) :
218 cc = type(self.ccfg)({}) # CostModelConfig({})
219 cc.parser = self
220 cc.config_format = "json"
221 cc.model_name = mod.model_id
222 cc.freeze = mod.freeze # for later
223 cc.has_fa = True
224 cc.has_op = True # mod.use_distributed_optimizer
225 cc.has_grad_shard = True
226 # cc.vp_less_mem = False
227 cc.has_clip = False
228 cc.cp_algo = "colossalai_cp"
229 cc.gmm = mod.moe_grouped_gemm
230 cc.vocab_emb_dp = False
232 cc.offset = 0
233 # Parallel dimensions
234 self.__config_parse_json_parallelism(cc, mod)
236 cc.full_rec = mod.recompute_num_layers
237 cc.sel_rec = False
238 if mod.recompute_num_layers and isinstance(
239 mod.recompute_num_layers, int
240 ):
241 cc.full_rec = [mod.recompute_num_layers] * cc.p
242 if cc.vp > 1:
243 cc.full_rec = [cc.full_rec] * cc.vp
245 # Hyperparameters
246 self.__config_parse_json_hyperparameters(cc, mod)
248 # Microbatch infos
249 cc.b = self.config.tmp.mbs # MBS # Microbatch size
250 cc.m = cc.p # Number of microbatches
251 # if cc.m<=0 : logger.warning("num_micro is negative")
253 # MoE infos
254 self.__config_parse_json_moe(cc, mod)
255 self.config_dp_tp_exp(cc)
257 # FP byte storages
258 cc.bytes_p = self.ccfg.fp_bytes(mod.params_dtype) # parameters
259 cc.bytes_compute = 2
260 cc.bytes_softmax = (
261 4 if mod.attention_softmax_in_fp32 else 2
262 ) # softmax output
263 cc.bytes_grad = 4
264 cc.bytes_os = 4
265 cc.bytes_norm = 4
267 # Optimizer parallel factors
268 cc.os_max_shard = cc.d * cc.t
269 self.config_optimizer_shard(cc)
271 # Other factors
272 cc.shard_embed = cc.t * cc.d
273 cc.shard_output_activ = 1
274 cc.shard_recompute_input = 1
275 cc.s_fa = (
276 cc.s if not cc.has_fa else cc.s / cc.a
277 ) # flash attention factor [HYPOTHESIS]
278 cc.comm_d_non_exp = (
279 0
280 if ((cc.d == 1) or not cc.has_op)
281 else (2 if not cc.has_grad_shard else 3)
282 ) # data parallel comm factor
283 cc.comm_d_exp = (
284 0
285 if ((cc.d_exp == 1) or not cc.has_op)
286 else (2 if not cc.has_grad_shard else 3)
287 ) # data parallel comm factor
288 cc.comm_t = float(cc.t > 1) # tensor parallel comm factor
289 cc.comm_ep = float(
290 cc.ep > 1 or cc.n_exp > 1
291 ) # expert parallel comm factor
292 cc.comm_cp = float(cc.cp > 1) # context parallel comm factor
293 cc.gbs = cc.b * cc.d * cc.m
294 cc.n_mtp = mod.mtp_num_layers
295 # Recomputation
296 self.__config_parse_json_op_recompute(cc)
297 cc.layer_custom_config = [(cc.n_lay, None)]
298 # By default, 100% of layers use a unique custom config (if specified)
299 cc.overwrite_eval_functions = {}
300 return cc # mod_hook