« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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 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
20from hyper_parallel.auto_parallel.sapp_nd.nd.common.derive import derive
23class CostModelParserMindspeed(_CostModelParser):
24 """parser class for MindSpeed format"""
26 def parse(self):
27 self.__config_parse_json_multimodals()
29 # MindSpeed (multimodal)
30 def __config_parse_json_multimodals(self):
31 """MindSpeed format for multimodal"""
32 self.ccfg.device_capacity = Memory.from_gb(55) # 55 * 1024 * 1024 * 1024
33 self.ccfg.model_name = self.config.model_id
34 self.config_op_counts(self.ccfg)
35 # Assume it exists a hook module with the same name as model_name
36 self.ccfg.n_lay = 0 # SUM ALL
37 self.ccfg.pp_sched = "1f1b"
38 self.ccfg.p = self.config.tmp.pp # PP
39 self.ccfg.m = self.ccfg.p
40 self.ccfg.b = self.config.tmp.mbs
41 self.ccfg.d = self.config.tmp.dp # DP
42 self.ccfg.t = self.config.tmp.tp # TP
43 self.ccfg.os_max_shard = self.ccfg.d
44 self.ccfg.cp = self.config.tmp.cp # CP
45 self.ccfg.vp = self.config.tmp.vpp # VPP
46 self.ccfg.ep = self.config.tmp.ep # EP
47 ccfgs = self.__search_and_parse_mods_ccfg(self.ccfg.config)
48 self.ccfg.multimodal = len(ccfgs) > 1
49 if len(ccfgs) == 1:
50 self.__adopt_module(next(iter(ccfgs.values())))
51 if self.ccfg.multimodal:
52 if not self.ccfg.hooks_dict:
53 raise TypeError(
54 "Currently for multimodals, 'hook_cls' required for evaluator"
55 )
56 missing_hooks = set(ccfgs.keys()) - set(self.ccfg.hooks_dict.keys())
57 if missing_hooks:
58 raise TypeError(
59 f"Missing hooks for submodules {list(missing_hooks)}"
60 )
61 self.ccfg.mm_ccfgs = ccfgs
62 self.ccfg.mm_order = list(ccfgs.keys())
64 # Update each mod's offset and pp_partition
65 for m in self.ccfg.mm_ccfgs:
66 cc = self.ccfg.mm_ccfgs[m]
67 num_layer_per_stage = max(1, cc.n_lay // self.ccfg.p // self.ccfg.vp)
68 if cc.pp_partition:
69 # print("here",m,cc.pp_partition,num_layer_per_stage)
70 # Making sure the format is offset[vp][p] like in MF
71 if isinstance(cc.pp_partition[0], list):
72 cc.offset = [
73 [
74 cc.pp_partition[v_idx][idx] - num_layer_per_stage
75 for idx in range(self.ccfg.p)
76 ]
77 for v_idx in range(self.ccfg.vp)
78 ]
79 else:
80 cc.offset = [
81 [
82 cc.pp_partition[idx] // self.ccfg.vp
83 - num_layer_per_stage
84 for idx in range(self.ccfg.p)
85 ]
86 for v_idx in range(self.ccfg.vp)
87 ]
88 else:
89 self.__complete_unimodal_pp_plan(m, cc, num_layer_per_stage)
90 self.ccfg.overwrite_eval_functions = {}
92 # What a config keeps of its own when it takes its one module's parse.
93 _KEPT_FROM_CONFIG = ("config", "config_path", "parser", "hooks_dict", "device_capacity",
94 "multimodal", "mm_ccfgs", "mm_order")
96 def __adopt_module(self, module):
97 """Make a config of one module that module: its model, run and derived fields."""
98 for key, value in vars(module).items():
99 if key not in self._KEPT_FROM_CONFIG:
100 object.__setattr__(self.ccfg, key, value)
102 @staticmethod
103 def __recompute_count(full_rec):
104 """The layers a module recomputes on each of its stages, as its parse states them."""
105 while isinstance(full_rec, list):
106 full_rec = next((count for count in full_rec if count), 0)
107 return int(full_rec or 0)
109 def __completed_recompute(self, cc):
110 """The module's own recompute, on the stages its completed plan puts it on."""
111 own = self.__recompute_count(cc.full_rec)
112 per_stage = [[min(own, count) for count in chunk] for chunk in cc.pp_partition]
113 return per_stage if cc.vp > 1 else per_stage[0]
115 def __complete_unimodal_pp_plan(self, m, cc, num_layer_per_stage):
116 """Try to follow previous pp plan"""
117 stage_insert_idx, chunk_insert_idx = 0, 0
118 cc.offset = [
119 [-num_layer_per_stage for _ in range(self.ccfg.p)]
120 for _ in range(self.ccfg.vp)
121 ]
122 previous_mod_idx = (
123 self.ccfg.mm_order.index(m) - 1
124 ) # look for previous mod pp partition
125 if previous_mod_idx >= 0:
126 previous_mod_partition = self.ccfg.mm_ccfgs[
127 self.ccfg.mm_order[previous_mod_idx]
128 ].pp_partition
129 if isinstance(previous_mod_partition[0], list): # vpp
130 put = False
131 for v_idx in range(self.ccfg.vp - 1, -1, -1):
132 for s_idx in range(self.ccfg.p - 1, -1, -1):
133 if previous_mod_partition[v_idx][s_idx]:
134 stage_insert_idx = s_idx
135 chunk_insert_idx = v_idx
136 put = True
137 break
138 if put:
139 break
140 else:
141 stage_insert_idx = self.ccfg.p - 1
142 for p in previous_mod_partition[::-1]:
143 if p > 0:
144 break
145 stage_insert_idx = max(0, stage_insert_idx - 1)
146 cc.offset[chunk_insert_idx][stage_insert_idx] = (
147 cc.n_lay - num_layer_per_stage
148 )
149 cc.pp_partition = [
150 [
151 num_layer_per_stage + cc.offset[v_idx][idx]
152 for idx in range(self.ccfg.p)
153 ]
154 for v_idx in range(self.ccfg.vp)
155 ]
156 cc.p, cc.vp = self.ccfg.p, self.ccfg.vp
157 cc.full_rec = self.__completed_recompute(cc)
158 cc.sel_rec = False
160 def __search_and_parse_mods_ccfg(self, field):
161 """extract multimodal submodules (MindSpeed format)"""
162 res = {}
163 for _, v in vars(field).items():
164 if isinstance(v, YamlObject):
165 res.update(self.__search_and_parse_mods_ccfg(v))
166 if v.model_id:
167 logger.info("Detected model config: %s", v.model_id)
168 res[v.model_id] = self.__config_parse_json(
169 v
170 ) # Build cost model variable foreach configs
171 return res
173 def __config_parse_json_parallelism(self, cc, mod):
174 """MindSpeed format for parallelism"""
175 cc.t = max(mod.tensor_model_parallel_size, self.config.tmp.tp)
176 cc.p = max(mod.pipeline_model_parallel_size, self.config.tmp.pp)
177 cc.cp = self.config.tmp.cp
178 cc.d = self.config.tmp.dp
179 cc.ep = max(mod.expert_model_parallel_size, self.config.tmp.ep)
180 cc.sequence_parallel = bool(mod.sequence_parallel)
181 if cc.cp > 1 and cc.sequence_parallel and cc.t > 1:
182 logger.warning(
183 "sequence parallelism and context parallelism are both enabled"
184 )
185 cc.pp_partition = (
186 mod.pipeline_num_layers
187 ) # Offset regardless of even distribution
189 # Interleaving
190 cc.n_s_split = 1 # seqpipe
191 cc.pp_sched = "1f1b"
192 cc.vp = self.config.tmp.vpp # VPP
194 def __config_parse_json_hyperparameters(self, cc, mod):
195 """MindSpeed format for hyperparams"""
196 cc.n_lay = mod.num_layers
197 cc.h = mod.hidden_size
198 cc.hff = mod.ffn_hidden_size
199 cc.v = mod.vocab_size
200 cc.s = self.config.tmp.seqlen # SEQ_LEN
201 cc.a = mod.num_attention_heads
202 cc.n_kv = (
203 mod.num_query_groups if mod.num_query_groups else cc.a
204 ) # NOT SURE
205 cc.dh = (
206 mod.kv_channels if mod.kv_channels else (cc.h / cc.a)
207 ) # Per head dimension
208 # print(cc.a, cc.h, cc.dh)
209 cc.dc_kv = mod.k_lora_rank # KV compression dimension #NOT SURE
210 cc.dc_q = mod.q_lora_rank # Q compression dimension
211 cc.dhr = mod.qk_rope_head_dim # decoupled QK per head dimension
212 cc.qk_norm = bool(mod.qk_layernorm) # Megatron's --qk-layernorm
214 def __config_parse_json_moe(self, cc, mod):
215 """MindSpeed format for MoE infos"""
216 cc.n_exp = max(1, mod.num_moe_experts)
217 cc.n_chosen_exp = max(1, mod.moe_router_topk)
218 cc.n_shared_exp = mod.n_shared_exp
219 cc.cap_fact = 1
220 cc.t_exp, cc.d_exp = cc.t, cc.d
221 cc.hff_exp = (
222 mod.moe_intermediate_size if mod.moe_intermediate_size else cc.hff
223 )
224 cc.k_1st_dense = mod.first_k_dense_replace
225 # temporary
226 cc.etp = self.config.tmp.etp # ETP
228 def __config_parse_json(self, mod):
229 """MindSpeed format for unimodal"""
230 # def mod_hook(M) :
231 cc = type(self.ccfg)({}) # CostModelConfig({})
232 cc.parser = self
233 cc.config_format = "json"
234 cc.model_name = mod.model_id
235 self.config_op_counts(cc)
236 cc.freeze = mod.freeze # for later
237 cc.has_fa = True
238 cc.has_op = True # mod.use_distributed_optimizer
239 cc.has_grad_shard = True
240 # cc.vp_less_mem = False
241 cc.has_clip = False
242 cc.cp_algo = "colossalai_cp"
243 cc.gmm = mod.moe_grouped_gemm
244 cc.vocab_emb_dp = False
245 cc.emb_dp_sharded = True
247 cc.offset = 0
248 # Parallel dimensions
249 self.__config_parse_json_parallelism(cc, mod)
251 cc.full_rec = mod.recompute_num_layers
252 cc.sel_rec = False # selective recompute is not parsed
253 cc.sel_rec_rule = "hyperparallel"
254 if mod.recompute_num_layers and isinstance(
255 mod.recompute_num_layers, int
256 ):
257 cc.full_rec = [mod.recompute_num_layers] * cc.p
258 if cc.vp > 1:
259 cc.full_rec = [cc.full_rec] * cc.vp
261 # Hyperparameters
262 self.__config_parse_json_hyperparameters(cc, mod)
264 # Microbatch infos
265 cc.b = self.config.tmp.mbs # MBS # Microbatch size
266 cc.m = cc.p # Number of microbatches
267 # if cc.m<=0 : logger.warning("num_micro is negative")
269 # MoE infos
270 self.__config_parse_json_moe(cc, mod)
272 # FP byte storages
273 cc.bytes_p = self.ccfg.fp_bytes(mod.params_dtype) # parameters
274 cc.bytes_compute = 2
275 cc.bytes_softmax = (
276 4 if mod.attention_softmax_in_fp32 else 2
277 ) # softmax output
279 # Optimizer parallel factors
280 cc.os_max_shard = cc.d
282 cc.gbs = cc.b * cc.d * cc.m
283 cc.n_mtp = mod.mtp_num_layers
284 self.config_layer_stack(cc)
285 cc.overwrite_eval_functions = {}
286 derive(cc)
287 return cc # mod_hook