« 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"""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 device_capacity: Memory = Memory.zero() # float = 0
43 mm_ccfgs: any = None
44 mm_order: list = None
45 layer_custom_config: list = None
46 overwrite_eval_functions: dict = None
47 parser: any = None
49 # Strategy
50 d: float = 0
51 t: float = 0
52 p: float = 0
53 cp: float = 0
54 ep: float = 1
55 sp: float = 0
56 vp: float = 0
57 os_max_shard: float = 0
58 op_weight_shard: float = 0
59 offset: Union[list, int] = None
60 full_rec: Union[list, bool] = None
61 sel_rec: Union[list, bool] = None
62 pp_sched: str = None
63 n_s_split: float = 0
64 cp_algo: str = "colossalai_cp"
65 rec_op: any = None
66 pp_partition: list = None
68 # hyperparameters
69 h: float = 0
70 hff: float = 0
71 v: float = 0
72 s: float = 0
73 s_fa: float = 0
74 a: float = 0
75 n_lay: float = 0
76 n_kv: float = 0
77 dh: float = 0
78 dc_kv: float = 0
79 dc_q: float = 0
80 dhr: float = 0
81 k_1st_dense: float = 0
82 n_mtp: float = 0
83 is_mtp_in_offset: bool = True
84 multiple_of: float = 0
85 fdm: float = 0
87 # MoE
88 t_exp: float = 0
89 d_exp: float = 0
90 hff_exp: float = 0
91 n_exp: float = 0
92 n_chosen_exp: float = 0
93 n_shared_exp: float = 0
94 cap_fact: float = 0
95 etp: float = 0
96 tokens_per_expert: list = None # global per-expert token count per microbatch (all EP ranks combined, before all-to-all); None = balanced
98 # CP modeling
99 kv_lora_rank: float = 0
100 attention_type: str = None
101 device_per_node: float = 8
102 bw_intra: float = 400.0
103 bw_inter: float = 25.0
104 sp_enabled: bool = False
106 # ZeRO
107 shard_p_os_non_exp_partial: float = 0
108 shard_p_os_non_exp: float = 0
109 shard_grad_non_exp: float = 0
110 shard_p_os_exp_partial: float = 0
111 shard_p_os_exp: float = 0
112 shard_grad_exp: float = 0
113 shard_grad_exp_partial: float = 0
115 # comm flag
116 comm_d_non_exp: float = 0
117 comm_d_exp: float = 0
118 comm_t: float = 0
119 comm_ep: float = 0
120 comm_cp: float = 0
121 # Transitional comm overlap correction (see comm_time.py).
122 # Fraction of comm volume hidden behind compute, applied as
123 # comm[dim] *= (1 - overlap_dim)
124 # Applies to both FLOP and TIME paths; the TIME path's
125 # estimate_comm_score(overlap=...) call site is zeroed so this is the
126 # single source of overlap.
127 # Defaults (dp=0.9, tp=0.5) are validated on MindFormers: they are the
128 # overlap that made the model match real MindFormers step times. Raw
129 # volume over-counts because communication really does overlap with
130 # compute. Re-validating for the hyper-parallel target is a follow-up.
131 # Follow-up: source from hardware profiling, then fold into
132 # estimate_comm_score (which also needs DP dedup, EP/CP terms, latency).
133 comm_dp_overlap: float = 0.9
134 comm_tp_overlap: float = 0.5
136 # feature flag
137 has_op: bool = False
138 has_grad_shard: bool = False
139 freeze: bool = False
140 has_fa: bool = False
141 # vp_less_mem: bool = False
142 has_clip: bool = False
143 gmm: bool = False
144 vocab_emb_dp: float = 0
145 tie_emb_out: bool = False
146 emb_out_in_offset: bool = False
148 # batch
149 b: float = 0
150 m: float = 0
151 gbs: float = 0
153 # shard
154 shard_embed: float = 0
155 shard_output_activ: float = 0
156 shard_recompute_input: float = 0
157 is_shard_mtp_param: bool = True
159 # bytes
160 bytes_p: float = 0
161 bytes_compute: float = 0
162 bytes_softmax: float = 0
163 bytes_grad: float = 0
164 bytes_os: float = 0
165 bytes_norm: float = 0
167 def __init__(self, input_config: Any, hook_cls: Any, framework: Optional[str], source_code: Optional[str]) -> None:
168 """Initialise from a config path and optional hooks/framework/source."""
169 super().__init__()
170 if input_config:
171 self.update_config(input_config, hook_cls, framework, source_code)
173 def _load_parser_cls(self, module_name):
174 """hook_class in eval yaml"""
175 target_mod_path = None
176 try:
177 # search in folder 'framework_parsers'
178 fram_dir = os.path.join(current_dir, "framework_parsers")
179 for f in os.listdir(fram_dir):
180 if f.endswith(".py"):
181 mod_path = f"hyper_parallel.auto_parallel.sapp_nd.nd.common.framework_parsers.{f.split('.')[0]}"
182 spec = importlib.util.find_spec(mod_path)
183 if spec is None or spec.origin is None:
184 continue
185 with open(spec.origin, "r", encoding="utf-8") as mf:
186 source = mf.read()
187 tree = ast.parse(source)
188 mod_cls = None
189 for node in ast.walk(tree):
190 if isinstance(node, ast.ClassDef) and node.name == module_name:
191 mod_cls = node
192 break
193 if mod_cls:
194 target_mod_path = mod_path
195 break
196 if target_mod_path:
197 module = importlib.import_module(target_mod_path)
198 return getattr(module, module_name)
199 except (ModuleNotFoundError, ImportError) as e:
200 print(e)
201 return None
203 def get_framework_parser_naive(self, input_config: str) -> Optional[Any]:
204 """Return parser class based on file extension (naive heuristic)."""
205 mod_name = None
206 if isinstance(input_config, str):
207 if input_config.endswith("yaml"):
208 mod_name = "CostModelParserMindformers"
209 if input_config.endswith("json"):
210 mod_name = "CostModelParserMindspeed"
211 if input_config.endswith("toml"):
212 mod_name = "CostModelParserHyperparallel"
213 if not mod_name:
214 raise AttributeError(f"Unhandled input format '{input_config}'")
215 return self._load_parser_cls(mod_name)
216 return None
218 def get_framework_parser(self, framework: str) -> Any:
219 """Look up and return the parser class for the given framework name.
221 Uses the mapping YAML file to find the corresponding parser module
222 class name, then loads and returns it. Raises AttributeError if the
223 framework name is not found in the mapping.
224 """
225 yml = Config(MAPPING_YML)
226 mod_name = next((e["module"] for e in yml.framework_parser if e["name"] == framework), None)
227 if not mod_name:
228 raise AttributeError(f"Cannot find parser module name from arg '{framework}'")
229 return self._load_parser_cls(mod_name)
231 def update_config(
232 self,
233 input_config: Any,
234 hook_cls: Any = None,
235 framework: Optional[str] = None,
236 source_code: Optional[str] = None,
237 ) -> None:
238 """process input config"""
239 self.hooks_dict = None if not hook_cls else hook_cls.get_hooks()
240 self.source_code = source_code
241 if isinstance(input_config, str):
242 self.config = Config(input_config)
243 # get parser
244 if framework:
245 logger.debug("Find parser module based on input framework name")
246 parser_cls = self.get_framework_parser(framework.lower())
247 else:
248 logger.debug("Naive way to find parser module")
249 parser_cls = self.get_framework_parser_naive(input_config)
250 if parser_cls:
251 self.parser = parser_cls(self)
252 logger.debug("Parser module: %s", self.parser.__class__)
253 self.parser.parse()
254 return
255 if isinstance(input_config, dict):
256 self.config = Config(input_config)
257 elif isinstance(input_config, Config):
258 self.config = input_config
259 else:
260 raise TypeError(
261 f"Expecting path string or Config object for {input_config}"
262 )
263 #MindFormers format by default
264 self.parser = self.get_framework_parser_naive("yaml")(self)
265 self.parser.parse()