« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 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 inspect
17import os
18from collections import Counter, defaultdict
19from contextlib import contextmanager
20from typing import Any, DefaultDict, Iterator, List
21import re
22from copy import deepcopy
23from enum import Enum
24from pprint import pformat
26from hyper_parallel.auto_parallel.sapp_nd.nd.common.generate_partitions import PartitionGenerator
27from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.logger import logger
28from hyper_parallel.auto_parallel.sapp_nd.nd.common.apply_exec import apply_exec, strategy_exec
31class AttentionType(Enum):
32 """Attention type enumeration."""
33 MHA = "mha"
34 GQA = "gqa"
35 MLA = "mla"
38def detect_attention_type(ccfg: "CostModelConfig") -> AttentionType:
39 """Detect attention type from cost model config.
41 Detection rules:
42 1. If the keys and values come from a compressed latent (dc_kv > 0,
43 the latent width every parser states): MLA
44 2. If n_kv < a: GQA
45 3. Otherwise: MHA
47 Args:
48 ccfg: Cost model config.
50 Returns:
51 AttentionType enum.
53 Example:
54 >>> ccfg.dc_kv = 512
55 >>> ccfg.a = 64
56 >>> ccfg.n_kv = 64
57 >>> detect_attention_type(ccfg)
58 <AttentionType.MLA: 'mla'>
59 """
60 if ccfg.dc_kv > 0:
61 return AttentionType.MLA
62 if ccfg.n_kv < ccfg.a:
63 return AttentionType.GQA
64 return AttentionType.MHA
67def compute_kv_dim(ccfg: Any) -> float:
68 """Return the width, per TP rank, of each of the K and V that CP exchanges.
70 Context parallelism exchanges the attention's own keys and values,
71 after the projections, whose heads TP splits. An MLA layer builds every
72 head's key and value from its latent before attention, the key at the
73 head width plus the rotary part every head shares, the value at the
74 head width, and those are what CP exchanges, not the latent: it acts on
75 the attention's inputs. Their mean width is returned, since K and V
76 are counted as two tensors of one width.
78 Args:
79 ccfg: Cost model config with attributes a, n_kv, dh, dhr, h, t, dc_kv.
81 Returns:
82 Effective KV dimension per TP rank (float).
83 """
84 attention_type = detect_attention_type(ccfg)
85 t = max(1, ccfg.t)
86 if attention_type == AttentionType.MLA:
87 return ccfg.n_kv * (2 * ccfg.dh + ccfg.dhr) / 2 / t
88 if attention_type == AttentionType.GQA:
89 n_kv = min(ccfg.n_kv if ccfg.n_kv > 0 else ccfg.a, ccfg.a)
90 return n_kv * ccfg.dh / t
91 return ccfg.h / t
94# Each read of a config field no parser set, by field and by the function
95# that read it. The config hands such a field over as 0, so a model feature
96# nobody parsed costs nothing unless its reader has a fallback of its own, and
97# the estimate stays plausible; a search says which were read, under its
98# ranking. A reader with a fallback, such as _flavour_tables' three
99# feed-forward matmuls, is named too: the record cannot tell it apart.
100UNSET_READS: DefaultDict[str, Counter] = defaultdict(Counter)
101# The fields whose absence means there is nothing to price: a submodule
102# copied without its parent's hooks takes the predefined ones.
103ABSENT_MEANS_NOTHING = frozenset({"hooks_dict"})
104# Whether such a read is refused rather than priced as 0 (run_nd --strict).
105_STRICT = {"on": False}
108def set_strict(strict: bool) -> None:
109 """Refuse, rather than price as 0, a config field no parser set where the estimate reads it.
111 Under it a read such as ``ccfg.n_gather`` raises AttributeError naming
112 the field and its reader, while ``getattr(ccfg, name, default)``, which
113 states what absence means, gets its default.
115 Args:
116 strict: True to refuse; False to price an unset field as 0, as by default.
117 """
118 _STRICT["on"] = bool(strict)
121def unset_reads_report() -> List[str]:
122 """What the config fields no parser set were read by, one line, and forget them; none where none was read."""
123 if not UNSET_READS:
124 return []
125 count = len(UNSET_READS)
126 fields = "; ".join(
127 f"{attr} by {', '.join(reader for reader, _ in readers.most_common())}"
128 for attr, readers in sorted(UNSET_READS.items())
129 )
130 UNSET_READS.clear()
131 return [f"The estimate read {count} config field(s) no parser set, each handed over as 0, which a reader "
132 f"without a fallback of its own prices as 0 (run_nd --strict refuses them): {fields}"]
135# class CostModelConfig(Config) :
136# The strategy fields a config armed by a search takes only through
137# set_strategy or apply_exec.
138STRATEGY_GUARDED = frozenset(("d", "t", "ep", "p", "vp", "cp", "os_max_shard"))
139# Where a config keeps whether its guard is armed; a copy never takes it.
140_GUARD = "_strategy_guard"
141# What a deep copy shares with its original: the parser and the raw yaml, set
142# when the config is parsed and written by no estimate. A search deep-copies
143# its config twice a candidate, and copying these was most of a copy's time.
144_SHARED_ON_COPY = frozenset(("parser", "config"))
147class CostModelConfig(PartitionGenerator):
148 """cost model variables class"""
150 def __init__(
151 self,
152 input_config: Any = None,
153 hook_cls: Any = None,
154 framework: Any = None,
155 source_code: Any = None,
156 devices: int = 0,
157 ) -> None:
158 """Initialise the cost model from a config, hooks and framework name, and the devices a caller states."""
159 super().__init__(input_config, hook_cls, framework, source_code, devices)
160 logger.debug(
161 "parser = %s for %s", str(self.parser), str(self.model_name)
162 )
164 def __str__(self) -> str:
165 """Return the cost model variables as readable text."""
166 return "CostModelConfig attributes:\n" + pformat(
167 {
168 k: v
169 for k, v in vars(self).items()
170 if isinstance(v, (int, float, str, bool))
171 }
172 )
174 def __getattr__(self, attr):
175 call_source = inspect.currentframe().f_back.f_code.co_name
176 if attr not in self.__dict__:
177 if not attr.startswith("_") and attr not in ABSENT_MEANS_NOTHING:
178 frame = inspect.currentframe().f_back
179 reader = f"{os.path.basename(frame.f_code.co_filename)}:{frame.f_code.co_name}"
180 UNSET_READS[attr][reader] += 1
181 if _STRICT["on"]:
182 raise AttributeError(
183 f"{reader} reads {attr}, which no parser set: run_nd --strict refuses it rather than "
184 f"pricing it as 0"
185 )
186 logger.warning(
187 "[%s] Attribute %s does not exist. "
188 "Value '0' will be assigned.",
189 call_source,
190 attr,
191 )
192 return 0
193 return self.__dict__[attr]
195 def __setattr__(self, name: str, value: Any) -> None:
196 if name in STRATEGY_GUARDED and self.__dict__.get(_GUARD):
197 raise AttributeError(f"Cannot directly modify {name}, use set_strategy()")
198 object.__setattr__(self, name, value)
200 def __copy__(self):
201 res = object.__new__(type(self))
202 res.__dict__.update(self.__dict__)
203 res.__dict__.pop(_GUARD, None)
204 return res
206 def __deepcopy__(self, memo):
207 res = object.__new__(type(self))
208 for k, v in self.__dict__.items():
209 if k != _GUARD:
210 setattr(res, k, v if k in _SHARED_ON_COPY else deepcopy(v, memo))
211 return res
213 def __getstate__(self) -> dict:
214 """Return instance state for multiprocessing serialization."""
215 state = self.__dict__.copy()
216 state.pop(_GUARD, None)
217 return state
219 def __setstate__(self, state: dict) -> None:
220 """Restore instance state after multiprocessing deserialization."""
221 self.__dict__.update(state)
223 def fp_bytes(self, precision: Any) -> int:
224 """Return bytes size for datatype"""
225 if precision and isinstance(precision, str):
226 res = re.match(r"[^0-9]*([0-9]+)[^0-9]*", precision)
227 if res:
228 return int(res.group(1)) // 8
229 logger.warning("No bytes detected from FP Precision: %s", precision)
230 return 0
232 def print_stages_i(self, stage_id: Any, stage: Any) -> None:
233 """for print_stages"""
234 stage_layers = []
235 for chunk in stage:
236 chunk_lay_occ = []
237 if chunk:
238 layer, count = chunk[0], 1
239 for lay_id in range(1, len(chunk)):
240 if chunk[lay_id] == layer:
241 count += 1
242 else:
243 chunk_lay_occ += [f"{count}{layer.name[0]}"]
244 layer, count = chunk[lay_id], 1
245 chunk_lay_occ += [f"{count}{layer.name[0]}"]
246 stage_layers += [chunk_lay_occ]
247 logger.info("stage _%s : %s", stage_id, stage_layers)
249 def print_stages(self, stages: Any, spec_stage_id: Any = -1) -> None:
250 """Call after generate_partitions"""
251 if spec_stage_id == -1:
252 for stage_id, stage in enumerate(stages):
253 self.print_stages_i(stage_id, stage)
254 elif 0 <= spec_stage_id < len(stages):
255 self.print_stages_i(spec_stage_id, stages[spec_stage_id])
256 else:
257 logger.warning("Incorrect spec_stage_id")
259 def count_layers(self, stages: Any) -> int:
260 """Count non-embedding and non-output layers in generated stages."""
261 return sum(sum(len(layer) for layer in chunk) for chunk in stages) - 2
263 def print_parallelism(self) -> None:
264 """strategy pretty printer"""
265 if not self.multimodal:
266 logger.info("%s Parallelism used :", self.model_name)
267 logger.info(
268 "DP %s, TP %s, PP %s, EP %s, CP %s, VPP %s",
269 self.d,
270 self.t,
271 self.p,
272 self.ep,
273 self.cp,
274 self.vp,
275 )
276 logger.info(
277 "d_exp %s, t_exp %s, os_max_shard %s, etp %s",
278 self.d_exp,
279 self.t_exp,
280 self.os_max_shard,
281 self.etp,
282 )
283 logger.info(
284 "shard_grad_exp %s, shard_grad_non_exp %s",
285 self.shard_grad_exp,
286 self.shard_grad_non_exp,
287 )
288 logger.info(
289 "shard_p_os_exp %s, shard_p_os_non_exp %s",
290 self.shard_p_os_exp,
291 self.shard_p_os_non_exp,
292 )
293 logger.info(
294 "shard_embed %s, shard_output_activ %s, shard_rec_input %s",
295 self.shard_embed,
296 self.shard_output_activ,
297 self.shard_recompute_input,
298 )
299 else:
300 for m in self.mm_ccfgs:
301 self.mm_ccfgs[m].print_parallelism()
303 def strategy_num_devices(self) -> float:
304 """total num devices"""
305 return self.d * self.t * self.cp * self.p
307 def is_consistent_pp_config(self) -> bool:
308 """check if pp/offset/recomputation consistency"""
310 def is_valid_cfg(cfg: Any) -> bool:
311 """Return whether one per-stage config matches the pipeline shape."""
312 if cfg is None or isinstance(cfg, (int, bool)):
313 return True
314 if not isinstance(cfg, list) or not cfg:
315 return False
316 if isinstance(cfg[0], int):
317 return len(cfg) == self.p
318 if isinstance(cfg[0], list):
319 return len(cfg) == self.vp and all(
320 isinstance(c, list) and len(c) == self.p for c in cfg
321 )
322 return False
324 return (
325 is_valid_cfg(self.offset)
326 and is_valid_cfg(self.full_rec)
327 and is_valid_cfg(self.sel_rec)
328 )
330 def __strategy_target(self, model_name):
331 """Get the config object targeted by a strategy update."""
332 if not self.multimodal:
333 return self
334 if model_name in self.mm_ccfgs:
335 return self.mm_ccfgs[model_name]
336 raise TypeError(
337 f"{self.model_name}: model_name is required (multimodal)"
338 )
340 def set_shared_strategy(self, **kwargs: Any) -> None:
341 """Apply a strategy to this config itself, a multimodal parent included.
343 :meth:`set_strategy` targets a submodule of a multimodal config,
344 never the parent; but the parent holds the strategy its submodules
345 share, and the partitions it combines are of that pipeline, so a
346 caller that gives every config of one pipeline its own placement
347 writes each of them through this.
348 """
349 apply_exec(self, strategy_exec(self, kwargs))
351 def set_strategy(self, **kwargs: Any) -> None:
352 """Apply a keyword strategy: the ExecSpec it states, and what follows.
354 The keywords are ``strategy_exec``'s; ``model_name`` picks the
355 submodule of a multimodal config, and without it every submodule
356 takes the strategy.
357 """
358 model_name = kwargs.get("model_name", None)
359 if self.multimodal and model_name is None:
360 # Submodules share one pipeline, so a strategy update with no
361 # explicit target has to reach every one of them.
362 for sub_name in self.mm_order:
363 self.mm_ccfgs[sub_name].set_strategy(
364 **{**kwargs, "model_name": None}
365 )
366 model_name = self.mm_main if self.mm_main else self.mm_order[-1]
367 target_ccfg = self.__strategy_target(model_name)
368 apply_exec(target_ccfg, strategy_exec(target_ccfg, kwargs))
369 logger.debug(
370 "in ccfg: DP = %d, TP = %d, EP = %d, CP = %d, "
371 "PP = %d, MB = %d, MBS = %d, VPP = %d",
372 target_ccfg.d,
373 target_ccfg.t,
374 target_ccfg.ep,
375 target_ccfg.cp,
376 target_ccfg.p,
377 target_ccfg.m,
378 target_ccfg.b,
379 target_ccfg.vp,
380 )
381 if not target_ccfg.is_consistent_pp_config():
382 raise AttributeError(
383 f"{target_ccfg.model_name}: "
384 "Inconsistent pipeline parallel variables "
385 f"pp {target_ccfg.p} vpp {target_ccfg.vp} "
386 f"offset {target_ccfg.offset} "
387 f"full_rec {target_ccfg.full_rec} "
388 f"sel_rec {target_ccfg.sel_rec}"
389 )
391 def get_strategy(self) -> Any:
392 """return parallelism/recompute strategies"""
394 def strategy(mm: Any) -> Any:
395 """Return the strategy dict for one submodule."""
396 return {
397 "dp": mm.d,
398 "tp": mm.t,
399 "pp": mm.p,
400 "ep": mm.ep,
401 "cp": mm.cp,
402 "vpp": mm.vp,
403 "op": mm.os_max_shard,
404 "gbs": mm.b * mm.m * mm.d,
405 "sched": mm.pp_sched,
406 "offset": mm.offset,
407 "full_rec": mm.full_rec,
408 "sel_rec": mm.sel_rec,
409 }
411 # logger.output("get_strat ccfg")
412 if self.multimodal:
413 return {mm.model_name: strategy(mm) for mm in self.mm_ccfgs.values()}
414 return strategy(self)
417def arm_strategy_guard(ccfg: Any, armed: bool = True) -> None:
418 """Refuse, or allow again, a direct write of a strategy field.
420 Armed, the config and its submodules take a new degree only through
421 ``set_strategy`` or ``apply_exec``, and a layer kind's own degrees
422 through ``apply_layer_strategy``. A copy of an armed config starts
423 unarmed, so an estimator can change its own.
425 Args:
426 ccfg: The config, multimodal or not.
427 armed: Whether to arm the guard or lift it.
428 """
429 object.__setattr__(ccfg, _GUARD, armed)
430 for sub in (getattr(ccfg, "mm_ccfgs", None) or {}).values():
431 arm_strategy_guard(sub, armed)
434@contextmanager
435def strategy_guarded(ccfg: Any) -> Iterator[Any]:
436 """Refuse a direct write of a strategy field on *ccfg* inside the block.
438 Afterwards the config is armed exactly as it was before.
439 """
440 before = ccfg.__dict__.get(_GUARD)
441 object.__setattr__(ccfg, _GUARD, True)
442 try:
443 yield ccfg
444 finally:
445 if before is None:
446 ccfg.__dict__.pop(_GUARD, None)
447 else:
448 object.__setattr__(ccfg, _GUARD, before)