« prev ^ index » next coverage.py v7.13.1, created at 2026-08-04 05:18 +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"""parse config for cost model"""
16import re
17import inspect
18from copy import deepcopy
19from pprint import pformat
20from enum import Enum
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
26class AttentionType(Enum):
27 """Attention type enumeration."""
28 MHA = "mha"
29 GQA = "gqa"
30 MLA = "mla"
33def detect_attention_type(ccfg: "CostModelConfig") -> AttentionType:
34 """Detect attention type from cost model config.
36 Detection rules:
37 1. If kv_lora_rank > 0: MLA
38 2. If n_kv < a: GQA
39 3. Otherwise: MHA
41 Args:
42 ccfg: Cost model config.
44 Returns:
45 AttentionType enum.
47 Example:
48 >>> ccfg.kv_lora_rank = 512
49 >>> ccfg.a = 64
50 >>> ccfg.n_kv = 64
51 >>> detect_attention_type(ccfg)
52 <AttentionType.MLA: 'mla'>
53 """
54 if ccfg.kv_lora_rank > 0:
55 return AttentionType.MLA
56 if ccfg.n_kv < ccfg.a:
57 return AttentionType.GQA
58 return AttentionType.MHA
61def compute_kv_dim(ccfg) -> float:
62 """Return effective KV dimension per TP rank based on attention type.
64 When TP is active, KV heads are split across TP ranks, so each
65 rank holds only 1/t of the total KV dimension. MLA is an
66 exception: the compressed latent vector is not split by TP,
67 so kv_lora_rank stays unchanged.
69 Args:
70 ccfg: Cost model config with attributes a, n_kv, dh, h, t, kv_lora_rank.
72 Returns:
73 Effective KV dimension per TP rank (float).
74 """
75 attention_type = detect_attention_type(ccfg)
76 t = max(1, ccfg.t)
77 if attention_type == AttentionType.MLA:
78 return float(ccfg.kv_lora_rank)
79 if attention_type == AttentionType.GQA:
80 n_kv = min(ccfg.n_kv if ccfg.n_kv > 0 else ccfg.a, ccfg.a)
81 return n_kv * ccfg.dh / t
82 return ccfg.h / t
85# class CostModelConfig(Config) :
86class CostModelConfig(PartitionGenerator):
87 """cost model variables class"""
89 def __init__(
90 self,
91 input_config=None,
92 hook_cls=None,
93 framework=None,
94 source_code=None,
95 ):
96 super().__init__(input_config, hook_cls, framework, source_code)
97 logger.debug(
98 "parser = %s for %s", str(self.parser), str(self.model_name)
99 )
101 def __str__(self):
102 return "CostModelConfig attributes:\n" + pformat(
103 {
104 k: v
105 for k, v in vars(self).items()
106 if isinstance(v, (int, float, str, bool))
107 }
108 )
110 def __getattr__(self, attr):
111 call_source = inspect.currentframe().f_back.f_code.co_name
112 if attr not in self.__dict__:
113 logger.warning(
114 "[%s] Attribute %s does not exist. "
115 "Value '0' will be assigned.",
116 call_source,
117 attr,
118 )
119 return 0
120 return self.__dict__[attr]
122 def __copy__(self):
123 res = object.__new__(type(self))
124 res.__dict__.update(self.__dict__)
125 return res
127 def __deepcopy__(self, memo):
128 res = object.__new__(type(self))
129 for k, v in self.__dict__.items():
130 setattr(res, k, deepcopy(v, memo))
131 return res
133 def fp_bytes(self, precision):
134 """Return bytes size for datatype"""
135 if precision and isinstance(precision, str):
136 res = re.match(r"[^0-9]*([0-9]+)[^0-9]*", precision)
137 if res:
138 return int(res.group(1)) // 8
139 logger.warning("No bytes detected from FP Precision: %s", precision)
140 return 0
142 def print_stages_i(self, stage_id, stage):
143 """for print_stages"""
144 stage_layers = []
145 for chunk in stage:
146 chunk_lay_occ = []
147 if chunk:
148 layer, count = chunk[0], 1
149 for lay_id in range(1, len(chunk)):
150 if chunk[lay_id] == layer:
151 count += 1
152 else:
153 chunk_lay_occ += [f"{count}{layer.name[0]}"]
154 layer, count = chunk[lay_id], 1
155 chunk_lay_occ += [f"{count}{layer.name[0]}"]
156 stage_layers += [chunk_lay_occ]
157 logger.info("stage _%s : %s", stage_id, stage_layers)
159 def print_stages(self, stages, spec_stage_id=-1):
160 """Call after generate_partitions"""
161 if spec_stage_id == -1:
162 for stage_id, stage in enumerate(stages):
163 self.print_stages_i(stage_id, stage)
164 elif 0 <= spec_stage_id < len(stages):
165 self.print_stages_i(spec_stage_id, stages[spec_stage_id])
166 else:
167 logger.warning("Incorrect spec_stage_id")
169 def count_layers(self, stages):
170 """Count non-embedding and non-output layers in generated stages."""
171 return sum(sum(len(layer) for layer in chunk) for chunk in stages) - 2
173 def print_parallelism(self):
174 """strategy pretty printer"""
175 if not self.multimodal:
176 logger.info("%s Parallelism used :", self.model_name)
177 logger.info(
178 "DP %s, TP %s, PP %s, EP %s, CP %s, VPP %s",
179 self.d,
180 self.t,
181 self.p,
182 self.ep,
183 self.cp,
184 self.vp,
185 )
186 logger.info(
187 "d_exp %s, t_exp %s, os_max_shard %s, etp %s",
188 self.d_exp,
189 self.t_exp,
190 self.os_max_shard,
191 self.etp,
192 )
193 logger.info(
194 "shard_grad_exp %s, shard_grad_non_exp %s",
195 self.shard_grad_exp,
196 self.shard_grad_non_exp,
197 )
198 logger.info(
199 "shard_p_os_exp %s, shard_p_os_non_exp %s",
200 self.shard_p_os_exp,
201 self.shard_p_os_non_exp,
202 )
203 logger.info(
204 "shard_embed %s, shard_output_activ %s, shard_rec_input %s",
205 self.shard_embed,
206 self.shard_output_activ,
207 self.shard_recompute_input,
208 )
209 else:
210 for m in self.mm_ccfgs:
211 self.mm_ccfgs[m].print_parallelism()
213 def strategy_num_devices(self):
214 """total num devices"""
215 return self.d * self.t * self.cp * self.p
217 def is_consistent_pp_config(self):
218 """check if pp/offset/recomputation consistency"""
220 def is_valid_cfg(cfg):
221 if cfg is None or isinstance(cfg, (int, bool)):
222 return True
223 if not isinstance(cfg, list) or not cfg:
224 return False
225 if isinstance(cfg[0], int):
226 return len(cfg) == self.p
227 if isinstance(cfg[0], list):
228 return len(cfg) == self.vp and all(
229 isinstance(c, list) and len(c) == self.p for c in cfg
230 )
231 return False
233 return (
234 is_valid_cfg(self.offset)
235 and is_valid_cfg(self.full_rec)
236 and is_valid_cfg(self.sel_rec)
237 )
239 @staticmethod
240 def __maybe_set_int(target, attr, value):
241 """Set an integer strategy attribute when an override is supplied."""
242 if isinstance(value, int):
243 setattr(target, attr, value)
245 def __strategy_target(self, model_name):
246 """Get the config object targeted by a strategy update."""
247 if not self.multimodal:
248 return self
249 if model_name in self.mm_ccfgs:
250 return self.mm_ccfgs[model_name]
251 raise TypeError(
252 f"{self.model_name}: model_name is required (multimodal)"
253 )
255 def set_strategy(self, **kwargs):
256 """overwrite parallelism"""
257 model_name = kwargs.get("model_name", None)
258 dp = kwargs.get("dp", None)
259 tp = kwargs.get("mp", None)
260 cp = kwargs.get("cp", None)
261 ep = kwargs.get("ep", None)
262 op = kwargs.get("op", None)
263 etp = kwargs.get("etp", None)
264 pp = kwargs.get("pp", None)
265 vpp = kwargs.get("vpp", None)
266 off = kwargs.get("offset", None)
267 fr = kwargs.get("full_rec", None)
268 sr = kwargs.get("sel_rec", None)
269 m = kwargs.get("mb", None)
270 b = kwargs.get("mbs", None)
271 target_ccfg = self.__strategy_target(model_name)
273 for attr, value in (
274 ("d", dp),
275 ("t", tp),
276 ("ep", ep),
277 ("etp", etp),
278 ("cp", cp),
279 ("vp", vpp),
280 ("p", pp),
281 ("m", m),
282 ("b", b),
283 ):
284 self.__maybe_set_int(target_ccfg, attr, value)
285 target_ccfg.sp = target_ccfg.t
286 if op is not None and isinstance(op, int):
287 target_ccfg.os_max_shard = op
288 # Sync has_op with os_max_shard: op<=1 means no optimizer sharding
289 target_ccfg.has_op = op > 1
290 target_ccfg.gbs = target_ccfg.b * target_ccfg.d * target_ccfg.m
291 logger.debug(
292 "in ccfg: DP = %d, TP = %d, EP = %d, CP = %d, "
293 "PP = %d, MB = %d, MBS = %d, VPP = %d",
294 target_ccfg.d,
295 target_ccfg.t,
296 target_ccfg.ep,
297 target_ccfg.cp,
298 target_ccfg.p,
299 target_ccfg.m,
300 target_ccfg.b,
301 target_ccfg.vp,
302 )
303 if hasattr(target_ccfg.parser, "config_shard_emb"):
304 target_ccfg.parser.config_shard_emb()
305 if hasattr(target_ccfg.parser, "config_shard_recompute"):
306 target_ccfg.parser.config_shard_recompute()
307 target_ccfg.parser.config_dp_tp_exp(target_ccfg)
308 target_ccfg.parser.config_optimizer_shard(target_ccfg)
309 target_ccfg.parser.config_comm_flag(target_ccfg)
310 if fr is not None:
311 target_ccfg.full_rec = fr
312 if sr is not None:
313 target_ccfg.sel_rec = sr
314 if isinstance(off, (int, list)):
315 target_ccfg.offset = off
316 if not target_ccfg.is_consistent_pp_config():
317 raise AttributeError(
318 f"{target_ccfg.model_name}: "
319 "Inconsistent pipeline parallel variables "
320 f"pp {target_ccfg.p} vpp {target_ccfg.vp} "
321 f"offset {target_ccfg.offset} "
322 f"full_rec {target_ccfg.full_rec} "
323 f"sel_rec {target_ccfg.sel_rec}"
324 )
325 self.__maybe_set_int(target_ccfg, "cp", cp)
327 def get_strategy(self):
328 """return parallelism/recompute strategies"""
330 def strategy(mm):
331 return {
332 "dp": mm.d,
333 "tp": mm.t,
334 "pp": mm.p,
335 "ep": mm.ep,
336 "cp": mm.cp,
337 "vpp": mm.vp,
338 "op": mm.os_max_shard,
339 "gbs": mm.b * mm.m * mm.d,
340 "sched": mm.pp_sched,
341 "offset": mm.offset,
342 "full_rec": mm.full_rec,
343 "sel_rec": mm.sel_rec,
344 }
346 # logger.output("get_strat ccfg")
347 if self.multimodal:
348 return {mm.model_name: strategy(mm) for mm in self.mm_ccfgs.values()}
349 return strategy(self)
351 def layer_custom_config_callback(self, fun):
352 """
353 Use input fun as callback for layer_custom_config
354 Only for overwriting cost model variables
355 """
356 for idx, f in enumerate(self.layer_custom_config):
358 def wrap(e, hook=f[1]):
359 hook(e)
360 if isinstance(e, CostModelConfig):
361 fun(self)
362 else:
363 e.set_ccfg(fun)
365 wrap.__name__ = f"{f[1].__name__}_{fun.__name__}"
366 self.layer_custom_config[idx] = (f[0], wrap)