« 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"""cost model variables"""
16from __future__ import annotations
17from typing import TYPE_CHECKING
19import importlib
20import ast
21import os
22from dataclasses import dataclass
23from hyper_parallel.auto_parallel.sapp_nd.nd.common.config import Config
24from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.size import Memory
25from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.logger import logger
27if TYPE_CHECKING:
28 from typing import Any, Optional, Union
30current_dir = os.path.dirname(os.path.abspath(__file__))
31MAPPING_YML = os.path.join(current_dir, "framework_parsers/mapping.yaml")
34@dataclass
35class _CostModVar:
36 """cost model variables class"""
38 config: any = None
39 config_format: str = None
40 multimodal: bool = False
41 model_name: str = None
42 # Op profile the model is priced with, and its op counts per layer kind;
43 # both settled by the parser, never matched from model_name.
44 arch: str = None
45 op_counts: dict = None
46 # A vision tower's arch is the vision profile; this is its language
47 # model's family, whose activation sharding it takes.
48 inherited_arch: str = None
49 device_capacity: Memory = Memory.zero() # float = 0
50 mm_ccfgs: any = None
51 mm_order: list = None
52 # The layer stack the parser settled (an auto_parallel LayerStack), and
53 # the fields each of its kinds assigns, bound per candidate by the arch
54 # hook from the config it is applied to.
55 layer_stack: any = None
56 layer_binding: dict = None
57 # What a layer of each kind keeps and holds per token, as the model
58 # spec's census states it (an auto_parallel KindActivations per kind
59 # name), and the record of the kind of the layer priced, which the
60 # parser or the arch hook binds; with none, the formulas price it.
61 census: dict = None
62 kind_activations: any = None
63 # The output layer's record, as the census states it.
64 output_census: any = None
65 overwrite_eval_functions: dict = None
66 parser: any = None
67 # The device count a caller states, as run_nd -d does, which a parser
68 # takes where its config cannot state one; 0 where none is stated.
69 devices: int = 0
70 # False where the config states no training sequence length and the
71 # parser costs the model's context limit in its place.
72 seq_len_stated: bool = True
73 # The MTP layers the model states and the AutoModels trainer does not
74 # build, so that nothing prices them; 0 where none was dropped.
75 mtp_unpriced: int = 0
76 # The activation checkpoint mode the run trains with, the trainer's
77 # default included, where a HyperParallel parser read one; None elsewhere.
78 stated_ac_mode: str = None
80 # Strategy
81 d: float = 0
82 t: float = 0
83 p: float = 0
84 cp: float = 0
85 ep: float = 1
86 # Whether activations are split along the sequence over the TP group,
87 # as the run's parser or a searched candidate states it; derive sets the
88 # factor sp from it. None where nothing stated it: priced as off, and a
89 # searched candidate then runs it at every TP, as the search always did.
90 sequence_parallel: bool = None
91 # Whether TP shards the activations between layers and the output
92 # layer's; None takes the family's (derive).
93 shard_activations: bool = None
94 # Whether the loss runs on logits sharded over the vocabulary, as stated,
95 # None taking its family's; and as derive gives it.
96 loss_parallel: bool = None
97 shards_logits: bool = True
98 sp: float = 0
99 vp: float = 0
100 os_max_shard: float = 0
101 op_weight_shard: float = 0
102 # How many ranks of its expert data-parallel group FSDP shards a routed
103 # expert over under expert parallelism, as the run states it; None for
104 # the optimizer's whole group.
105 expert_shard: int = None
106 # Whether the run shards each strategy's routed experts over that
107 # strategy's whole expert data-parallel group, as a launcher that sets
108 # edp_shard_size to the group does: the shard then follows the strategy.
109 expert_shard_group: bool = False
110 offset: Union[list, int] = None
111 full_rec: Union[list, bool] = None
112 sel_rec: Union[list, bool] = None
113 # MindFormers' select_comm_recompute, which its selective recompute reads.
114 sel_comm_rec: Union[list, bool] = None
115 # The framework whose selective recompute rec_op follows, "hyperparallel"
116 # or "mindformers".
117 sel_rec_rule: str = "hyperparallel"
118 # Whether a recomputed layer keeps its input sliced over tensor parallelism.
119 recompute_slice_activation: bool = False
120 # How each layer recomputes, as ranges over the layers in model order
121 # (ExecSpec.recompute), when a spec states them; they replace full_rec
122 # and sel_rec, and the partition generator places them.
123 recompute_ranges: tuple = None
124 pp_sched: str = None
125 n_s_split: float = 0
126 cp_algo: str = "colossalai_cp"
127 rec_op: any = None
128 # Each layer's own recompute switches in model order, None for a layer
129 # that is not selective, where the recompute ranges state several
130 # selective settings; None where rec_op holds the one every selective
131 # layer runs (derive).
132 layer_switches: tuple = None
133 pp_partition: list = None
135 # hyperparameters
136 h: float = 0
137 hff: float = 0
138 v: float = 0
139 s: float = 0
140 s_fa: float = 0
141 a: float = 0
142 n_lay: float = 0
143 n_kv: float = 0
144 dh: float = 0
145 # An MLA model's value-head width, at which its family prices dh (derive).
146 v_head_dim: float = None
147 # An MLA model's non-rotary key-head width, its value heads' unless stated.
148 qk_nope_head_dim: float = None
149 dc_kv: float = 0
150 dc_q: float = 0
151 dhr: float = 0
152 # The biases and norms a model states (its spec's); None where unstated,
153 # which the parameter formulas count their own way.
154 qkv_bias: bool = None
155 o_bias: bool = None
156 mlp_bias: bool = None
157 norm_bias: bool = None
158 layer_norms: int = None
159 shared_expert_gate: bool = None
160 k_1st_dense: float = 0
161 # Attention flavour of a layer group, and the extra parameters a flavour
162 # carries that the q/k/v/o formula does not describe (conv, gates).
163 attn_kind: str = "full"
164 attn_extra_p: float = 0
165 lin_n_k: float = 0
166 lin_d_k: float = 0
167 lin_n_v: float = 0
168 lin_d_v: float = 0
169 lin_conv: float = 0
170 # Recurrent-state update and readout, the linear-attention op. Zero for
171 # every other flavour.
172 n_linrec: float = 0
173 # The QK-norm a layer runs: 1 where the model normalizes each head's
174 # queries and keys (derive, from qk_norm), 0 on a linear-attention layer.
175 n_qknorm: float = 0
176 n_mtp: float = 0
177 is_mtp_in_offset: bool = True
178 multiple_of: float = 0
179 fdm: float = 0
181 # MoE
182 t_exp: float = 0
183 d_exp: float = 0
184 # One per MoE layer: the op count of the dispatch the layer runs to put
185 # each token on the rank holding the expert it routed to, and to take the
186 # result back (``_flavour_tables``).
187 n_dispatch: float = 0
188 # What that dispatch costs on this cluster, as ``context.moe_dispatch``
189 # states it: None, unstated, keeps the measured default
190 # (``estimate.MOE_DISPATCH``), and 0 prices no dispatch at all.
191 moe_dispatch: float = None
192 # The ranks that hold the same experts, as the runtime's expert mesh spans
193 # them; see ``_CostModelParser.expert_dp_group``. ``d_exp`` is the older
194 # per-CP, per-expert-TP-rank form of it and cannot carry CP.
195 edp_group: float = 0
196 hff_exp: float = 0
197 n_exp: float = 0
198 n_chosen_exp: float = 0
199 n_shared_exp: float = 0
200 cap_fact: float = 0
201 etp: float = 0
202 # Global per-expert token count per microbatch (all EP ranks combined,
203 # before all-to-all); None = balanced.
204 tokens_per_expert: list = None
206 # CP modeling
207 attention_type: str = None
208 device_per_node: float = 8
209 bw_intra: float = 400.0
210 bw_inter: float = 25.0
211 sp_enabled: bool = False
213 # ZeRO
214 shard_p_os_non_exp_partial: float = 0
215 shard_p_os_non_exp: float = 0
216 shard_grad_non_exp: float = 0
217 shard_p_os_exp_partial: float = 0
218 shard_p_os_exp: float = 0
219 shard_grad_exp: float = 0
220 shard_grad_exp_partial: float = 0
222 # comm flag
223 comm_d_non_exp: float = 0
224 comm_d_exp: float = 0
225 comm_t: float = 0
226 comm_ep: float = 0
227 comm_cp: float = 0
228 # Transitional comm overlap correction (see comm_time.py).
229 # Fraction of comm volume hidden behind compute, applied as
230 # comm[dim] *= (1 - overlap_dim)
231 # Applies to both FLOP and TIME paths; the TIME path's
232 # estimate_comm_score(overlap=...) call site is zeroed so this is the
233 # single source of overlap.
234 # Defaults (dp=0.9, tp=0.5) are validated on MindFormers: they are the
235 # overlap that made the model match real MindFormers step times. Raw
236 # volume over-counts because communication really does overlap with
237 # compute. Re-validating for the hyper-parallel target is a follow-up.
238 # Follow-up: source from hardware profiling, then fold into
239 # estimate_comm_score (which also needs DP dedup, EP/CP terms, latency).
240 comm_dp_overlap: float = 0.9
241 comm_tp_overlap: float = 0.5
243 # feature flag
244 has_op: bool = False
245 has_grad_shard: bool = False
246 # Whether each gradient is sharded as its parameter is, as FSDP holds it.
247 grad_shard_as_params: bool = False
248 # Whether the stated FSDP shard spans CP's ranks, as HyperParallel's
249 # dp_shard_size does, so that CP shards no parameter state beyond it
250 # (derive.param_cp); False keeps the legacy schemas' rule.
251 shard_spans_cp: bool = False
252 freeze: bool = False
253 has_fa: bool = False
254 attn_output_gate: bool = False
255 # Whether attention normalizes each head's queries and keys (Qwen3).
256 qk_norm: bool = False
257 # vp_less_mem: bool = False
258 has_clip: bool = False
259 gmm: bool = False
260 vocab_emb_dp: float = 0
261 # Whether the embedding table is sharded over data parallelism too.
262 emb_dp_sharded: bool = True
263 tie_emb_out: bool = False
264 emb_out_in_offset: bool = False
266 # batch
267 b: float = 0
268 m: float = 0
269 gbs: float = 0
271 # shard
272 shard_embed: float = 0
273 # Over how many of the ranks that shard the embedding table its layer
274 # gathers it to compute with it.
275 gather_embed: float = 1
276 shard_output_activ: float = 0
277 shard_recompute_input: float = 0
278 is_shard_mtp_param: bool = True
280 # bytes
281 bytes_p: float = 0
282 bytes_compute: float = 0
283 bytes_softmax: float = 0
284 # The widths the run states, None taking its family's, and whether its
285 # gradients take memory without pipeline parallelism; derive gives the
286 # estimators bytes_grad, bytes_os, bytes_norm and bytes_dropout from them.
287 grad_bytes: float = None
288 optimizer_state_bytes: float = None
289 optimizer_states: float = None
290 main_param_bytes: float = None
291 norm_bytes: float = None
292 dropout_bytes: float = None
293 grad_accumulation: bool = None
294 # Whether the run accumulates gradients over micro-batches without
295 # pipeline parallelism, as derive gives it from grad_accumulation: a
296 # search then gives PP 1 several micro-batches.
297 accumulates_grads: bool = False
298 # Whether FSDP frees a layer's gathered parameters once it has run, as
299 # the run states it, None taking its family's; and as derive gives it.
300 reshard_params: bool = None
301 reshards: bool = False
302 # Whether FSDP holds each layer's reduce-scatter output until the
303 # backward ends, as stated, None taking its family's; and as derive
304 # gives it.
305 deferred_grad_accumulation: bool = None
306 defers_grads: bool = False
307 # Whether FSDP holds a layer's whole gradients while the next layer's
308 # backward runs, and the root's until the backward ends, as stated,
309 # None taking its family's; and as derive gives it.
310 overlapped_grad_reduce: bool = None
311 overlaps_grad_reduce: bool = False
312 # Whether a layer keeps the weight casts the formulas price, as stated,
313 # None where the optimizer does not shard; and as derive gives it.
314 param_casts: bool = None
315 keeps_param_casts: bool = None
316 bytes_grad: float = 0
317 bytes_os: float = 0
318 # What the optimizer keeps per parameter: a layer's, and the embedding
319 # and output tables', which keep AdamW's two states (derive).
320 bytes_optim: float = 0
321 bytes_optim_table: float = 0
322 bytes_norm: float = 0
323 bytes_dropout: float = 0
325 def __init__(
326 self,
327 input_config: Any,
328 hook_cls: Any,
329 framework: Optional[str],
330 source_code: Optional[str],
331 devices: int = 0,
332 ) -> None:
333 """Initialise from a config path and optional hooks/framework/source, and the devices a caller states."""
334 super().__init__()
335 self.devices = devices or 0
336 if input_config:
337 self.update_config(input_config, hook_cls, framework, source_code)
339 def _load_parser_cls(self, module_name):
340 """hook_class in eval yaml"""
341 target_mod_path = None
342 try:
343 # search in folder 'framework_parsers'
344 fram_dir = os.path.join(current_dir, "framework_parsers")
345 for f in os.listdir(fram_dir):
346 if f.endswith(".py"):
347 mod_path = f"hyper_parallel.auto_parallel.sapp_nd.nd.common.framework_parsers.{f.split('.')[0]}"
348 spec = importlib.util.find_spec(mod_path)
349 if spec is None or spec.origin is None:
350 continue
351 with open(spec.origin, "r", encoding="utf-8") as mf:
352 source = mf.read()
353 tree = ast.parse(source)
354 mod_cls = None
355 for node in ast.walk(tree):
356 if isinstance(node, ast.ClassDef) and node.name == module_name:
357 mod_cls = node
358 break
359 if mod_cls:
360 target_mod_path = mod_path
361 break
362 if target_mod_path:
363 module = importlib.import_module(target_mod_path)
364 return getattr(module, module_name)
365 except (ModuleNotFoundError, ImportError) as e:
366 print(e)
367 return None
369 def get_framework_parser_naive(self, input_config: str) -> Optional[Any]:
370 """Return parser class based on file extension (naive heuristic)."""
371 mod_name = None
372 if isinstance(input_config, str):
373 if input_config.endswith("yaml"):
374 mod_name = "CostModelParserMindformers"
375 if input_config.endswith("json"):
376 mod_name = "CostModelParserMindspeed"
377 if input_config.endswith("toml"):
378 mod_name = "CostModelParserHyperparallel"
379 if not mod_name:
380 raise AttributeError(f"Unhandled input format '{input_config}'")
381 return self._load_parser_cls(mod_name)
382 return None
384 def get_framework_parser(self, framework: str) -> Any:
385 """Look up and return the parser class for the given framework name.
387 Uses the mapping YAML file to find the corresponding parser module
388 class name, then loads and returns it. Raises AttributeError if the
389 framework name is not found in the mapping.
390 """
391 yml = Config(MAPPING_YML)
392 mod_name = next((e["module"] for e in yml.framework_parser if e["name"] == framework), None)
393 if not mod_name:
394 raise AttributeError(f"Cannot find parser module name from arg '{framework}'")
395 return self._load_parser_cls(mod_name)
397 def update_config(
398 self,
399 input_config: Any,
400 hook_cls: Any = None,
401 framework: Optional[str] = None,
402 source_code: Optional[str] = None,
403 ) -> None:
404 """process input config"""
405 self.hooks_dict = None if not hook_cls else hook_cls.get_hooks()
406 self.source_code = source_code
407 if isinstance(input_config, (str, dict)):
408 self.config = Config(input_config)
409 elif isinstance(input_config, Config):
410 self.config = input_config
411 else:
412 raise TypeError(
413 f"Expecting path string, dict or Config object for {input_config}"
414 )
415 # An in-memory config names its framework the same way a file does,
416 # so a caller that already holds the config need not write it to disk
417 # for the naive path sniffer to find the parser.
418 if framework:
419 logger.debug("Find parser module based on input framework name")
420 parser_cls = self.get_framework_parser(framework.lower())
421 elif isinstance(input_config, str):
422 logger.debug("Naive way to find parser module")
423 parser_cls = self.get_framework_parser_naive(input_config)
424 else:
425 # MindFormers format by default
426 parser_cls = self.get_framework_parser_naive("yaml")
427 if parser_cls:
428 self.parser = parser_cls(self)
429 logger.debug("Parser module: %s", self.parser.__class__)
430 self.parser.parse()