« 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"""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 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 = 0
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: float = 0
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 # ZeRO
99 shard_p_os_non_exp_partial: float = 0
100 shard_p_os_non_exp: float = 0
101 shard_grad_non_exp: float = 0
102 shard_p_os_exp_partial: float = 0
103 shard_p_os_exp: float = 0
104 shard_grad_exp: float = 0
105 shard_grad_exp_partial: float = 0
107 # comm flag
108 comm_d_non_exp: float = 0
109 comm_d_exp: float = 0
110 comm_t: float = 0
111 comm_ep: float = 0
112 comm_cp: float = 0
114 # feature flag
115 has_op: bool = False
116 has_grad_shard: bool = False
117 freeze: bool = False
118 has_fa: bool = False
119 # vp_less_mem: bool = False
120 has_clip: bool = False
121 gmm: bool = False
122 vocab_emb_dp: float = 0
123 tie_emb_out: bool = False
124 emb_out_in_offset: bool = False
126 # batch
127 b: float = 0
128 m: float = 0
129 gbs: float = 0
131 # shard
132 shard_embed: float = 0
133 shard_output_activ: float = 0
134 shard_recompute_input: float = 0
135 is_shard_mtp_param: bool = True
137 # bytes
138 bytes_p: float = 0
139 bytes_compute: float = 0
140 bytes_softmax: float = 0
141 bytes_grad: float = 0
142 bytes_os: float = 0
143 bytes_norm: float = 0
145 def __init__(self, input_config, hook_cls, framework, source_code):
146 super().__init__()
147 if input_config:
148 self.update_config(input_config, hook_cls, framework, source_code)
150 def _load_parser_cls(self, module_name):
151 """hook_class in eval yaml"""
152 target_mod_path = None
153 try:
154 # search in folder 'framework_parsers'
155 fram_dir = os.path.join(current_dir, "framework_parsers")
156 for f in os.listdir(fram_dir):
157 if f.endswith(".py"):
158 mod_path = f"hyper_parallel.auto_parallel.sapp_nd.nd.common.framework_parsers.{f.split('.')[0]}"
159 spec = importlib.util.find_spec(mod_path)
160 if spec is None or spec.origin is None:
161 continue
162 with open(spec.origin, "r", encoding="utf-8") as mf:
163 source = mf.read()
164 tree = ast.parse(source)
165 mod_cls = None
166 for node in ast.walk(tree):
167 if isinstance(node, ast.ClassDef) and node.name == module_name:
168 mod_cls = node
169 break
170 if mod_cls:
171 target_mod_path = mod_path
172 break
173 if target_mod_path:
174 module = importlib.import_module(target_mod_path)
175 return getattr(module, module_name)
176 except (ModuleNotFoundError, ImportError) as e:
177 print(e)
178 return None
180 def get_framework_parser_naive(self, input_config):
181 "yaml for MindFormers, json for Mindspeed, toml for HyperParallel"
182 mod_name = None
183 if isinstance(input_config, str):
184 if input_config.endswith("yaml"):
185 mod_name = "CostModelParserMindformers"
186 if input_config.endswith("json"):
187 mod_name = "CostModelParserMindspeed"
188 if input_config.endswith("toml"):
189 mod_name = "CostModelParserHyperparallel"
190 if not mod_name:
191 raise AttributeError(f"Unhandled input format '{input_config}'")
192 return self._load_parser_cls(mod_name)
193 return None
195 def get_framework_parser(self, framework):
196 """Look up and return the parser class for the given framework name.
198 Uses the mapping YAML file to find the corresponding parser module
199 class name, then loads and returns it. Raises AttributeError if the
200 framework name is not found in the mapping.
201 """
202 yml = Config(MAPPING_YML)
203 mod_name = next((e["module"] for e in yml.framework_parser if e["name"] == framework), None)
204 if not mod_name:
205 raise AttributeError(f"Cannot find parser module name from arg '{framework}'")
206 return self._load_parser_cls(mod_name)
208 def update_config(self, input_config, hook_cls=None, framework=None, source_code=None):
209 """process input config"""
210 self.hooks_dict = None if not hook_cls else hook_cls.get_hooks()
211 self.source_code = source_code
212 if isinstance(input_config, str):
213 self.config = Config(input_config)
214 # get parser
215 if framework:
216 logger.debug("Find parser module based on input framework name")
217 parser_cls = self.get_framework_parser(framework.lower())
218 else:
219 logger.debug("Naive way to find parser module")
220 parser_cls = self.get_framework_parser_naive(input_config)
221 if parser_cls:
222 self.parser = parser_cls(self)
223 logger.debug("Parser module: %s", self.parser.__class__)
224 self.parser.parse()
225 return
226 if isinstance(input_config, dict):
227 self.config = Config(input_config)
228 elif isinstance(input_config, Config):
229 self.config = input_config
230 else:
231 raise TypeError(
232 f"Expecting path string or Config object for {input_config}"
233 )
234 #MindFormers format by default
235 self.parser = self.get_framework_parser_naive("yaml")(self)
236 self.parser.parse()