« 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"""hook manager module"""
16from __future__ import annotations
17from typing import TYPE_CHECKING
18import ast
19import textwrap
20import inspect
21from hyper_parallel.auto_parallel.sapp_nd.nd.common.config import Config
22from hyper_parallel.auto_parallel.sapp_nd.nd.common.layer_type import LayerType
23from hyper_parallel.auto_parallel.sapp_nd.nd.common.cost_model_preprocess import CostModelConfig
24from hyper_parallel.auto_parallel.sapp_nd.nd.common.arch_hooks import check_and_apply_custom_hook
25from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.logger import logger
26from hyper_parallel.auto_parallel.sapp_nd.memory_estimation._context import (
27 NodeEval,
28 NodeStatEval,
29 NodeDynEval,
30 NodeCommEval,
31)
32from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.head import EvalHead
33from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.tail import EvalTail
34from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.body import EvalBody
35from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.layer_block import EvalAttn, EvalFFn
36from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.layer_block import EvalNorm
37from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.comm import EvalLayerComm
38from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.utils import EvalUtils
39from hyper_parallel.auto_parallel.sapp_nd.memory_estimation._backbone import _Backbone
40from hyper_parallel.auto_parallel.sapp_nd.memory_estimation._func_tracer import _FuncTracer
41from hyper_parallel.auto_parallel.sapp_nd.memory_estimation._context import MemType
43if TYPE_CHECKING:
44 from typing import Callable, Any
47class _HookManager(_Backbone):
48 """Hook manager class."""
50 def __init__(self, *args, **kwargs):
51 super().__init__(*args, **kwargs)
52 self.func_tracer = _FuncTracer()
53 self.toggle_func_trace = kwargs.get("trace_fun", False)
54 self.import_eval_yaml()
55 self.fetch_hook_if_unimodal()
57 def fetch_hook_if_unimodal(self):
58 """Process custom model config"""
59 if not self._ccfg.multimodal:
60 if not self._ccfg.hooks_dict:
61 logger.info(
62 "'hook_cls' not specified,"
63 "search in predefined arch_hooks"
64 )
65 check_and_apply_custom_hook(self)
66 else:
67 hook = list(self._ccfg.hooks_dict.values())[0]
68 hook(self)
70 def __is_valid_eval_func(self, fun: Any) -> bool:
71 """check hook definition, return type"""
72 if fun is None:
73 return False
74 if isinstance(fun, (str, int, float)):
75 return True
77 source = inspect.getsource(fun)
78 tree = ast.parse(textwrap.dedent(source))
79 for instruction in ast.walk(tree):
80 if (
81 isinstance(instruction, ast.Return)
82 and instruction.value is None
83 ):
84 return False
85 return True
87 # Evaluation context setters
89 def set_passes(
90 self,
91 vpp_less_mem: bool = None,
92 swap_os: bool = None,
93 dropless_tok_factor: float = None,
94 ) -> None:
95 """toggle features"""
96 if isinstance(vpp_less_mem, bool):
97 self._ctx.vpp_less_mem = vpp_less_mem
98 if isinstance(swap_os, bool):
99 self._ctx.swap_os = swap_os
100 if isinstance(dropless_tok_factor, (int, float)):
101 self._ctx.dropless_tok_factor = dropless_tok_factor
103 def __set_node_eval_fun(self, cls_obj, target_node, *args, **kwargs):
104 """overwrite given node's formulas"""
105 c_stat, c_dyn = None, None
106 if (args and args[-1] == 0) or kwargs.get("stat", None) == 0:
107 c_stat = 0
108 if (args and args[-1] == 0) or kwargs.get("dyn", None) == 0:
109 c_dyn = 0
110 num_p = kwargs.get("num_p", c_stat)
111 stat_p = kwargs.get("stat_p", c_stat)
112 stat_os = kwargs.get("stat_os", c_stat)
113 stat_grad = kwargs.get("stat_grad", c_dyn)
114 dyn_activ = kwargs.get("dyn_activ", c_dyn)
115 if not self.__is_valid_eval_func(num_p):
116 num_p = self._ctx.node_eval[target_node].num_p
117 if not self.__is_valid_eval_func(stat_p):
118 stat_p = self._ctx.node_eval[target_node].stat.p
119 if not self.__is_valid_eval_func(stat_os):
120 stat_os = self._ctx.node_eval[target_node].stat.os
121 if not self.__is_valid_eval_func(stat_grad):
122 stat_grad = self._ctx.node_eval[target_node].stat.grad
123 if not self.__is_valid_eval_func(dyn_activ):
124 dyn_activ = self._ctx.node_eval[target_node].dyn.activation
125 self._ctx.node_eval[target_node] = NodeEval(
126 self.__custom_getattr(cls_obj, num_p),
127 NodeStatEval(
128 self.__custom_getattr(cls_obj, stat_p, MemType.MODEL_PARAM),
129 self.__custom_getattr(cls_obj, stat_os, MemType.OPTIM_STATE),
130 self.__custom_getattr(cls_obj, stat_grad, MemType.ACCU_GRAD),
131 ),
132 NodeDynEval(
133 self.__custom_getattr(cls_obj, dyn_activ),
134 self.__set_node_eval_comm_fun(
135 cls_obj, target_node, *args, **kwargs
136 ),
137 ),
138 )
140 def __set_node_eval_comm_fun(self, cls_obj, target_node, *args, **kwargs):
141 """overwrite given node's comm formulas"""
142 c_comm = None
143 last_arg_is_zero = args and args[-1] == 0
144 dyn_comm_is_zero = kwargs.get("dyn_comm") == 0
145 dyn_is_zero = kwargs.get("dyn") == 0
147 if last_arg_is_zero or dyn_comm_is_zero or dyn_is_zero:
148 c_comm = 0
149 dyn_dp_comm = kwargs.get("dyn_dp_comm", c_comm)
150 dyn_tp_comm = kwargs.get("dyn_tp_comm", c_comm)
151 dyn_cp_comm = kwargs.get("dyn_cp_comm", c_comm)
152 dyn_ep_comm = kwargs.get("dyn_ep_comm", c_comm)
153 dyn_ep_comm_balanced = kwargs.get("dyn_ep_comm_balanced", None)
154 dyn_ep_comm_imbalanced = kwargs.get("dyn_ep_comm_imbalanced", None)
155 if not self.__is_valid_eval_func(dyn_dp_comm):
156 dyn_dp_comm = self._ctx.node_eval[target_node].dyn.comm.dp
157 if not self.__is_valid_eval_func(dyn_tp_comm):
158 dyn_tp_comm = self._ctx.node_eval[target_node].dyn.comm.tp
159 if not self.__is_valid_eval_func(dyn_cp_comm):
160 dyn_cp_comm = self._ctx.node_eval[target_node].dyn.comm.cp
161 if not self.__is_valid_eval_func(dyn_ep_comm):
162 dyn_ep_comm = self._ctx.node_eval[target_node].dyn.comm.ep
163 comm_cls_obj = cls_obj
164 if self.is_regular_layer(target_node):
165 comm_cls_obj = EvalLayerComm
166 ep_balanced = (
167 self.__custom_getattr(comm_cls_obj, dyn_ep_comm_balanced)
168 if self.__is_valid_eval_func(dyn_ep_comm_balanced)
169 else None
170 )
171 ep_imbalanced = (
172 self.__custom_getattr(comm_cls_obj, dyn_ep_comm_imbalanced)
173 if self.__is_valid_eval_func(dyn_ep_comm_imbalanced)
174 else None
175 )
176 return NodeCommEval(
177 self.__custom_getattr(comm_cls_obj, dyn_dp_comm),
178 self.__custom_getattr(comm_cls_obj, dyn_tp_comm),
179 self.__custom_getattr(comm_cls_obj, dyn_cp_comm),
180 self.__custom_getattr(comm_cls_obj, dyn_ep_comm),
181 ep_balanced=ep_balanced,
182 ep_imbalanced=ep_imbalanced,
183 )
185 def set_head_eval_fun(self, *arg, **kwarg):
186 """overwrite head formulas"""
187 self.__set_node_eval_fun(EvalHead, self._ctx.head_node, *arg, **kwarg)
189 def set_tail_eval_fun(self, *arg, **kwarg):
190 """overwrite tail formulas"""
191 self.__set_node_eval_fun(EvalTail, self._ctx.tail_node, *arg, **kwarg)
193 def set_body_eval_fun(self, *args, **kwargs):
194 """overwrite body formulas"""
195 lay_type = kwargs.get("lay_type", args[0] if args else None)
196 if not lay_type:
197 lt = [
198 b_obj
199 for b_obj in list(LayerType)
200 if self.is_regular_layer(b_obj)
201 ]
202 else:
203 if not isinstance(lay_type, LayerType):
204 b_obj = self.__custom_getattr(LayerType, lay_type)
205 lt = [b_obj]
206 else:
207 lt = [lay_type]
208 for b_obj in lt:
209 self.__set_node_eval_fun(EvalBody, b_obj, *args, **kwargs)
211 def set_attn_eval_fun(
212 self,
213 num_p: Any = None,
214 qkv: Any = None,
215 score: Any = None,
216 proj: Any = None,
217 ) -> None:
218 """overwrite attention formulas"""
219 if self.__is_valid_eval_func(num_p):
220 self._ctx.attn_num_p = self.__custom_getattr(EvalAttn, num_p)
221 if self.__is_valid_eval_func(qkv):
222 self._ctx.attn_qkv_activ = self.__custom_getattr(
223 EvalAttn, qkv, MemType.ATTN_ACTIV
224 )
225 if self.__is_valid_eval_func(score):
226 self._ctx.attn_score_activ = self.__custom_getattr(
227 EvalAttn, score, MemType.ATTN_ACTIV
228 )
229 if self.__is_valid_eval_func(proj):
230 self._ctx.attn_proj_activ = self.__custom_getattr(
231 EvalAttn, proj, MemType.ATTN_ACTIV
232 )
234 def set_ffn_eval_fun(self, num_p: Any = None, activation=None, moe_activ=None):
235 """overwrite feedforward formulas"""
236 if self.__is_valid_eval_func(num_p):
237 self._ctx.ffn_num_p = self.__custom_getattr(EvalFFn, num_p)
238 if self.__is_valid_eval_func(activation):
239 self._ctx.ffn_activ = self.__custom_getattr(
240 EvalFFn, activation, MemType.FFN_ACTIV
241 )
242 if self.__is_valid_eval_func(moe_activ):
243 self._ctx.ffn_moe_activ = self.__custom_getattr(
244 EvalFFn, moe_activ, MemType.FFN_ACTIV
245 )
247 def set_expert_param_eval_fun(
248 self, routed_num_p: Any = None, shared_num_p: Any = None
249 ) -> None:
250 """overwrite expert param count formulas for routed/shared breakdown"""
251 if self.__is_valid_eval_func(routed_num_p):
252 self._ctx.ffn_routed_num_p = self.__custom_getattr(
253 EvalFFn, routed_num_p
254 )
255 if self.__is_valid_eval_func(shared_num_p):
256 self._ctx.ffn_shared_num_p = self.__custom_getattr(
257 EvalFFn, shared_num_p
258 )
260 def set_norm_eval_fun(self, num_p: Any = None, activation=None):
261 """overwrite norm formulas"""
262 if self.__is_valid_eval_func(num_p):
263 self._ctx.norm_num_p = self.__custom_getattr(EvalNorm, num_p)
264 if self.__is_valid_eval_func(activation):
265 self._ctx.norm_activ = self.__custom_getattr(
266 EvalNorm, activation, MemType.NORM_ACTIV
267 )
269 def set_pp_micro_factor_eval_fun(self, sched_name, fun):
270 """overwrite PP microfactor formulas"""
271 if sched_name and self.__is_valid_eval_func(fun):
272 self._ctx.pp_micro_eval[sched_name] = self.__custom_getattr(
273 EvalUtils, fun
274 )
276 # Cost Model Config setter
278 def set_strategy(self, **kwargs):
279 """overwrite parallelism"""
280 self._ccfg.set_strategy(**kwargs)
281 self.fetch_hook_if_unimodal()
283 def set_ccfg(self, hook):
284 """overwrite cost model variable (except strategy)"""
285 if hook and callable(hook):
287 def custom_setter(self, name, value):
288 strat_vars = ["d", "t", "ep", "p", "vp", "cp", "os_max_shard"]
289 if name in strat_vars:
290 raise AttributeError(
291 f"Cannot directly modify {name}, use set_strategy()"
292 )
293 self.__dict__[name] = value
295 CostModelConfig.__setattr__ = custom_setter
296 hook(self._ccfg)
297 CostModelConfig.__setattr__ = object.__setattr__
299 def __wrap_mem_counter(self, mem_type: MemType, fun: Callable) -> None:
300 """Wrap formula calls to accumulate memory by type."""
301 if mem_type and not hasattr(fun, "wrapped_with_counter"):
303 def wrap(*args, **kwargs):
304 res = fun(*args, **kwargs)
305 self._ctx.accu_mem_type[mem_type] += res
306 self._ctx.save2log(mem_type, res)
307 return res
309 wrap.__qualname__ = fun.__qualname__
310 wrap.wrapped_with_counter = True
311 return wrap
312 return fun
314 def __custom_getattr(
315 self, eval_class: Any, field: Any, mem_type: MemType = None
316 ) -> Callable:
317 """formula retrieve/wrap"""
318 # Definition priority order :
319 # 1. Callable (user defined in code)
320 # OR Numeric value
321 # 2. Overriding list from cost_model_preprocess (ccfg attribute)
322 # 3. Function name (config_eval yaml)
323 res = None
324 if callable(field):
325 res = field
326 if self.toggle_func_trace == field.__name__:
327 res = self.func_tracer.wrap(field)
328 if isinstance(field, (int, float)):
330 def constant(*_):
331 return field
333 constant.__qualname__ = str(field)
334 res = constant
335 if field in self._ccfg.overwrite_eval_functions:
336 res = self._ccfg.overwrite_eval_functions[field]
337 if isinstance(field, str):
339 def zero(*_):
340 return 0
342 zero.__qualname__ = "0"
343 res = getattr(eval_class, field, zero)
344 if self.toggle_func_trace == field:
345 res = self.func_tracer.wrap(getattr(eval_class, field))
346 res = self.__wrap_mem_counter(mem_type, res)
347 if not res:
348 raise TypeError(f"In eval config yaml, non valid field: {field}")
349 return res
351 # Import Eval Config
353 def import_eval_yaml(self) -> None:
354 """import evaluator config file, init ctx (inner call only)"""
355 if not self.toggle_func_trace:
356 self.toggle_func_trace = self.eval_cfg.trace_fun
357 # head
358 h_obj = getattr(LayerType, self.eval_cfg.nodes_mem_comp.head.name)
359 self._ctx.head_node = h_obj
360 self.set_head_eval_fun(
361 num_p=self.eval_cfg.nodes_mem_comp.head.num_param_fun,
362 stat_p=self.eval_cfg.nodes_mem_comp.head.stat_fun.p,
363 stat_os=self.eval_cfg.nodes_mem_comp.head.stat_fun.os,
364 stat_grad=self.eval_cfg.nodes_mem_comp.head.stat_fun.grad,
365 dyn_activ=self.eval_cfg.nodes_mem_comp.head.dyn_fun.activation,
366 dyn_dp_comm=self.eval_cfg.nodes_mem_comp.head.dyn_fun.comm.dp,
367 dyn_tp_comm=self.eval_cfg.nodes_mem_comp.head.dyn_fun.comm.tp,
368 dyn_cp_comm=self.eval_cfg.nodes_mem_comp.head.dyn_fun.comm.cp,
369 dyn_ep_comm=self.eval_cfg.nodes_mem_comp.head.dyn_fun.comm.ep,
370 )
372 # tail
373 t_obj = getattr(LayerType, self.eval_cfg.nodes_mem_comp.tail.name)
374 self._ctx.tail_node = t_obj
375 self.set_tail_eval_fun(
376 num_p=self.eval_cfg.nodes_mem_comp.tail.num_param_fun,
377 stat_p=self.eval_cfg.nodes_mem_comp.tail.stat_fun.p,
378 stat_os=self.eval_cfg.nodes_mem_comp.tail.stat_fun.os,
379 stat_grad=self.eval_cfg.nodes_mem_comp.tail.stat_fun.grad,
380 dyn_activ=self.eval_cfg.nodes_mem_comp.tail.dyn_fun.activation,
381 dyn_dp_comm=self.eval_cfg.nodes_mem_comp.tail.dyn_fun.comm.dp,
382 dyn_tp_comm=self.eval_cfg.nodes_mem_comp.tail.dyn_fun.comm.tp,
383 dyn_cp_comm=self.eval_cfg.nodes_mem_comp.tail.dyn_fun.comm.cp,
384 dyn_ep_comm=self.eval_cfg.nodes_mem_comp.tail.dyn_fun.comm.ep,
385 )
387 # body
388 for b in self.eval_cfg.nodes_mem_comp.body:
389 b_cfg = Config(b)
390 comm_cfg = b_cfg.dyn_fun.comm
391 comm_kwargs = {
392 "dyn_dp_comm": comm_cfg.dp,
393 "dyn_tp_comm": comm_cfg.tp,
394 "dyn_cp_comm": comm_cfg.cp,
395 "dyn_ep_comm": comm_cfg.ep,
396 }
397 if hasattr(comm_cfg, "ep_balanced"):
398 comm_kwargs["dyn_ep_comm_balanced"] = comm_cfg.ep_balanced
399 if hasattr(comm_cfg, "ep_imbalanced"):
400 comm_kwargs["dyn_ep_comm_imbalanced"] = comm_cfg.ep_imbalanced
401 self.set_body_eval_fun(
402 lay_type=b_cfg.name,
403 num_p=b_cfg.num_param_fun,
404 stat_p=b_cfg.stat_fun.p,
405 stat_os=b_cfg.stat_fun.os,
406 stat_grad=b_cfg.stat_fun.grad,
407 dyn_activ=b_cfg.dyn_fun.activation,
408 **comm_kwargs,
409 )
411 # pp micro factor
412 for sc in self.eval_cfg.pp_sched:
413 self.set_pp_micro_factor_eval_fun(sc["name"], sc["fun"])
414 if not self._ccfg.pp_sched:
415 self._ccfg.pp_sched = self.eval_cfg.default_pp_sched
417 # layerblock
418 self.set_attn_eval_fun(
419 self.eval_cfg.base_arch_mem_comp.attention.num_param_fun,
420 self.eval_cfg.base_arch_mem_comp.attention.qkv,
421 self.eval_cfg.base_arch_mem_comp.attention.score,
422 self.eval_cfg.base_arch_mem_comp.attention.proj,
423 )
424 self.set_ffn_eval_fun(
425 self.eval_cfg.base_arch_mem_comp.feedforward.num_param_fun,
426 self.eval_cfg.base_arch_mem_comp.feedforward.activation,
427 self.eval_cfg.base_arch_mem_comp.feedforward.moe_activ,
428 )
429 self.set_expert_param_eval_fun(
430 routed_num_p=self.eval_cfg.base_arch_mem_comp.feedforward.routed_num_fun,
431 shared_num_p=self.eval_cfg.base_arch_mem_comp.feedforward.shared_num_fun,
432 )
433 self.set_norm_eval_fun(
434 self.eval_cfg.base_arch_mem_comp.norm.num_param_fun,
435 self.eval_cfg.base_arch_mem_comp.norm.activation,
436 )
438 # passes
439 self.set_passes(
440 vpp_less_mem=self.eval_cfg.passes.vpp_less_memory,
441 swap_os=self.eval_cfg.passes.swap_optimizer,
442 dropless_tok_factor=self.eval_cfg.passes.dropless_tok_factor,
443 )
445 self._ctx.comm_expr = self.eval_cfg.comm_expr
447 def is_regular_layer(self, lay):
448 """check if layer is not head/tail"""
449 if isinstance(lay, str):
450 return lay[0] not in [
451 self._ctx.head_node.name[0],
452 self._ctx.tail_node.name[0],
453 ]
454 return lay not in [self._ctx.head_node, self._ctx.tail_node]