« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +0800
1# Copyright 2025 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"""parse config for cost model"""
16import re
17import inspect
18from copy import deepcopy
19from pprint import pformat
21# from hyper_parallel.auto_parallel.sapp_nd.nd.config import Config, YamlObject
22from hyper_parallel.auto_parallel.sapp_nd.nd.common.generate_partitions import PartitionGenerator
23from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.logger import logger
26# class CostModelConfig(Config) :
27class CostModelConfig(PartitionGenerator):
28 """cost model variables class"""
30 def __init__(
31 self,
32 input_config=None,
33 hook_cls=None,
34 framework=None,
35 source_code=None,
36 ):
37 super().__init__(input_config, hook_cls, framework, source_code)
38 logger.debug(
39 "parser = %s for %s", str(self.parser), str(self.model_name)
40 )
42 def __str__(self):
43 return "CostModelConfig attributes:\n" + pformat(
44 {
45 k: v
46 for k, v in vars(self).items()
47 if isinstance(v, (int, float, str, bool))
48 }
49 )
51 def __getattr__(self, attr):
52 call_source = inspect.currentframe().f_back.f_code.co_name
53 if attr not in self.__dict__:
54 logger.warning(
55 "[%s] Attribute %s does not exist. "
56 "Value '0' will be assigned.",
57 call_source,
58 attr,
59 )
60 return 0
61 return self.__dict__[attr]
63 def __copy__(self):
64 res = object.__new__(type(self))
65 res.__dict__.update(self.__dict__)
66 return res
68 def __deepcopy__(self, memo):
69 res = object.__new__(type(self))
70 for k, v in self.__dict__.items():
71 setattr(res, k, deepcopy(v, memo))
72 return res
74 def fp_bytes(self, precision):
75 """Return bytes size for datatype"""
76 if precision and isinstance(precision, str):
77 res = re.match(r"[^0-9]*([0-9]+)[^0-9]*", precision)
78 if res:
79 return int(res.group(1)) // 8
80 logger.warning("No bytes detected from FP Precision: %s", precision)
81 return 0
83 def print_stages_i(self, stage_id, stage):
84 """for print_stages"""
85 stage_layers = []
86 for chunk in stage:
87 chunk_lay_occ = []
88 if chunk:
89 layer, count = chunk[0], 1
90 for lay_id in range(1, len(chunk)):
91 if chunk[lay_id] == layer:
92 count += 1
93 else:
94 chunk_lay_occ += [f"{count}{layer.name[0]}"]
95 layer, count = chunk[lay_id], 1
96 chunk_lay_occ += [f"{count}{layer.name[0]}"]
97 stage_layers += [chunk_lay_occ]
98 logger.info("stage _%s : %s", stage_id, stage_layers)
100 def print_stages(self, stages, spec_stage_id=-1):
101 """Call after generate_partitions"""
102 if spec_stage_id == -1:
103 for stage_id, stage in enumerate(stages):
104 self.print_stages_i(stage_id, stage)
105 elif 0 <= spec_stage_id < len(stages):
106 self.print_stages_i(spec_stage_id, stages[spec_stage_id])
107 else:
108 logger.warning("Incorrect spec_stage_id")
110 def count_layers(self, stages):
111 """Count non-embedding and non-output layers in generated stages."""
112 return sum(sum(len(layer) for layer in chunk) for chunk in stages) - 2
114 def print_parallelism(self):
115 """strategy pretty printer"""
116 if not self.multimodal:
117 logger.info("%s Parallelism used :", self.model_name)
118 logger.info(
119 "DP %s, TP %s, PP %s, EP %s, CP %s, VPP %s",
120 self.d,
121 self.t,
122 self.p,
123 self.ep,
124 self.cp,
125 self.vp,
126 )
127 logger.info(
128 "d_exp %s, t_exp %s, os_max_shard %s, etp %s",
129 self.d_exp,
130 self.t_exp,
131 self.os_max_shard,
132 self.etp,
133 )
134 logger.info(
135 "shard_grad_exp %s, shard_grad_non_exp %s",
136 self.shard_grad_exp,
137 self.shard_grad_non_exp,
138 )
139 logger.info(
140 "shard_p_os_exp %s, shard_p_os_non_exp %s",
141 self.shard_p_os_exp,
142 self.shard_p_os_non_exp,
143 )
144 logger.info(
145 "shard_embed %s, shard_output_activ %s, shard_rec_input %s",
146 self.shard_embed,
147 self.shard_output_activ,
148 self.shard_recompute_input,
149 )
150 else:
151 for m in self.mm_ccfgs:
152 self.mm_ccfgs[m].print_parallelism()
154 def strategy_num_devices(self):
155 """total num devices"""
156 return self.d * self.t * self.cp * self.p
158 def is_consistent_pp_config(self):
159 """check if pp/offset/recomputation consistency"""
161 def is_valid_cfg(cfg):
162 if cfg is None or isinstance(cfg, (int, bool)):
163 return True
164 if not isinstance(cfg, list) or not cfg:
165 return False
166 if isinstance(cfg[0], int):
167 return len(cfg) == self.p
168 if isinstance(cfg[0], list):
169 return len(cfg) == self.vp and all(
170 isinstance(c, list) and len(c) == self.p for c in cfg
171 )
172 return False
174 return (
175 is_valid_cfg(self.offset)
176 and is_valid_cfg(self.full_rec)
177 and is_valid_cfg(self.sel_rec)
178 )
180 @staticmethod
181 def __maybe_set_int(target, attr, value):
182 """Set an integer strategy attribute when an override is supplied."""
183 if isinstance(value, int):
184 setattr(target, attr, value)
186 def __strategy_target(self, model_name):
187 """Get the config object targeted by a strategy update."""
188 if not self.multimodal:
189 return self
190 if model_name in self.mm_ccfgs:
191 return self.mm_ccfgs[model_name]
192 raise TypeError(
193 f"{self.model_name}: model_name is required (multimodal)"
194 )
196 def set_strategy(self, **kwargs):
197 """overwrite parallelism"""
198 model_name = kwargs.get("model_name", None)
199 dp = kwargs.get("dp", None)
200 tp = kwargs.get("mp", None)
201 cp = kwargs.get("cp", None)
202 ep = kwargs.get("ep", None)
203 op = kwargs.get("op", None)
204 etp = kwargs.get("etp", None)
205 pp = kwargs.get("pp", None)
206 vpp = kwargs.get("vpp", None)
207 off = kwargs.get("offset", None)
208 fr = kwargs.get("full_rec", None)
209 sr = kwargs.get("sel_rec", None)
210 m = kwargs.get("mb", None)
211 b = kwargs.get("mbs", None)
212 target_ccfg = self.__strategy_target(model_name)
214 for attr, value in (
215 ("d", dp),
216 ("t", tp),
217 ("ep", ep),
218 ("etp", etp),
219 ("cp", cp),
220 ("vp", vpp),
221 ("p", pp),
222 ("m", m),
223 ("b", b),
224 ):
225 self.__maybe_set_int(target_ccfg, attr, value)
226 target_ccfg.sp = target_ccfg.t
227 if op is not None and isinstance(op, int):
228 target_ccfg.os_max_shard = op
229 # Sync has_op with os_max_shard: op<=1 means no optimizer sharding
230 target_ccfg.has_op = op > 1
231 target_ccfg.gbs = target_ccfg.b * target_ccfg.d * target_ccfg.m
232 logger.debug(
233 "in ccfg: DP = %d, TP = %d, EP = %d, CP = %d, "
234 "PP = %d, MB = %d, MBS = %d, VPP = %d",
235 target_ccfg.d,
236 target_ccfg.t,
237 target_ccfg.ep,
238 target_ccfg.cp,
239 target_ccfg.p,
240 target_ccfg.m,
241 target_ccfg.b,
242 target_ccfg.vp,
243 )
244 if hasattr(target_ccfg.parser, "config_shard_emb"):
245 target_ccfg.parser.config_shard_emb()
246 target_ccfg.parser.config_dp_tp_exp(target_ccfg)
247 target_ccfg.parser.config_optimizer_shard(target_ccfg)
248 target_ccfg.parser.config_comm_flag(target_ccfg)
249 if fr is not None:
250 target_ccfg.full_rec = fr
251 if sr is not None:
252 target_ccfg.sel_rec = sr
253 if isinstance(off, (int, list)):
254 target_ccfg.offset = off
255 if not target_ccfg.is_consistent_pp_config():
256 raise AttributeError(
257 f"{target_ccfg.model_name}: "
258 "Inconsistent pipeline parallel variables "
259 f"pp {target_ccfg.p} vpp {target_ccfg.vp} "
260 f"offset {target_ccfg.offset} "
261 f"full_rec {target_ccfg.full_rec} "
262 f"sel_rec {target_ccfg.sel_rec}"
263 )
264 self.__maybe_set_int(target_ccfg, "cp", cp)
266 def get_strategy(self):
267 """return parallelism/recompute strategies"""
269 def strategy(mm):
270 return {
271 "dp": mm.d,
272 "tp": mm.t,
273 "pp": mm.p,
274 "ep": mm.ep,
275 "cp": mm.cp,
276 "vpp": mm.vp,
277 "op": mm.os_max_shard,
278 "gbs": mm.b * mm.m * mm.d,
279 "sched": mm.pp_sched,
280 "offset": mm.offset,
281 "full_rec": mm.full_rec,
282 "sel_rec": mm.sel_rec,
283 }
285 # logger.output("get_strat ccfg")
286 if self.multimodal:
287 return {mm.model_name: strategy(mm) for mm in self.mm_ccfgs.values()}
288 return strategy(self)
290 def layer_custom_config_callback(self, fun):
291 """
292 Use input fun as callback for layer_custom_config
293 Only for overwriting cost model variables
294 """
295 for idx, f in enumerate(self.layer_custom_config):
297 def wrap(e, hook=f[1]):
298 hook(e)
299 if isinstance(e, CostModelConfig):
300 fun(self)
301 else:
302 e.set_ccfg(fun)
304 wrap.__name__ = f"{f[1].__name__}_{fun.__name__}"
305 self.layer_custom_config[idx] = (f[0], wrap)