« 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"""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 strategy_guarded
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 NodeComputeEval,
32)
33from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.compute import EvalExpertCompute
34from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.head import EvalHead
35from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.tail import EvalTail
36from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.body import EvalBody
37from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.layer_block import EvalAttn, EvalFFn
38from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.layer_block import EvalNorm
39from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.comm import EvalLayerComm
40from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.utils import EvalUtils
41from hyper_parallel.auto_parallel.sapp_nd.memory_estimation._backbone import _Backbone
42from hyper_parallel.auto_parallel.sapp_nd.memory_estimation._func_tracer import _FuncTracer
43from hyper_parallel.auto_parallel.sapp_nd.memory_estimation._context import MemType
45if TYPE_CHECKING:
46 from typing import Callable, Any
49class _HookManager(_Backbone):
50 """Hook manager class."""
52 def __init__(self, *args, **kwargs):
53 super().__init__(*args, **kwargs)
54 self.func_tracer = _FuncTracer()
55 self.toggle_func_trace = kwargs.get("trace_fun", False)
56 self.import_eval_yaml()
57 self.fetch_hook_if_unimodal()
59 def fetch_hook_if_unimodal(self):
60 """Process custom model config"""
61 if not self._ccfg.multimodal:
62 if not self._ccfg.hooks_dict:
63 logger.info(
64 "'hook_cls' not specified,"
65 "search in predefined arch_hooks"
66 )
67 check_and_apply_custom_hook(self)
68 else:
69 hook = list(self._ccfg.hooks_dict.values())[0]
70 hook(self)
72 def __is_valid_eval_func(self, fun: Any) -> bool:
73 """check hook definition, return type"""
74 if fun is None:
75 return False
76 if isinstance(fun, (str, int, float)):
77 return True
79 source = inspect.getsource(fun)
80 tree = ast.parse(textwrap.dedent(source))
81 for instruction in ast.walk(tree):
82 if (
83 isinstance(instruction, ast.Return)
84 and instruction.value is None
85 ):
86 return False
87 return True
89 # Evaluation context setters
91 def set_passes(
92 self,
93 vpp_less_mem: bool = None,
94 swap_os: bool = None,
95 dropless_tok_factor: float = None,
96 ) -> None:
97 """toggle features"""
98 if isinstance(vpp_less_mem, bool):
99 self._ctx.vpp_less_mem = vpp_less_mem
100 if isinstance(swap_os, bool):
101 self._ctx.swap_os = swap_os
102 if isinstance(dropless_tok_factor, (int, float)):
103 self._ctx.dropless_tok_factor = dropless_tok_factor
105 def __set_node_eval_fun(self, cls_obj, target_node, *args, **kwargs):
106 """overwrite given node's formulas"""
107 c_stat, c_dyn = None, None
108 if (args and args[-1] == 0) or kwargs.get("stat", None) == 0:
109 c_stat = 0
110 if (args and args[-1] == 0) or kwargs.get("dyn", None) == 0:
111 c_dyn = 0
112 num_p = kwargs.get("num_p", c_stat)
113 stat_p = kwargs.get("stat_p", c_stat)
114 stat_os = kwargs.get("stat_os", c_stat)
115 stat_grad = kwargs.get("stat_grad", c_dyn)
116 dyn_activ = kwargs.get("dyn_activ", c_dyn)
117 if not self.__is_valid_eval_func(num_p):
118 num_p = self._ctx.node_eval[target_node].num_p
119 if not self.__is_valid_eval_func(stat_p):
120 stat_p = self._ctx.node_eval[target_node].stat.p
121 if not self.__is_valid_eval_func(stat_os):
122 stat_os = self._ctx.node_eval[target_node].stat.os
123 if not self.__is_valid_eval_func(stat_grad):
124 stat_grad = self._ctx.node_eval[target_node].stat.grad
125 if not self.__is_valid_eval_func(dyn_activ):
126 dyn_activ = self._ctx.node_eval[target_node].dyn.activation
127 # Compute formulas (FLOPs, not memory — bypass mem_counter)
128 compute_eval = self.__set_node_eval_compute_fun(target_node, **kwargs)
129 self._ctx.node_eval[target_node] = NodeEval(
130 self.__custom_getattr(cls_obj, num_p),
131 NodeStatEval(
132 self.__custom_getattr(cls_obj, stat_p, MemType.MODEL_PARAM),
133 self.__custom_getattr(cls_obj, stat_os, MemType.OPTIM_STATE),
134 self.__custom_getattr(cls_obj, stat_grad, MemType.ACCU_GRAD),
135 ),
136 NodeDynEval(
137 self.__custom_getattr(cls_obj, dyn_activ),
138 self.__set_node_eval_comm_fun(
139 cls_obj, target_node, *args, **kwargs
140 ),
141 compute=compute_eval,
142 ),
143 )
145 def __set_node_eval_comm_fun(self, cls_obj, target_node, *args, **kwargs):
146 """overwrite given node's comm formulas"""
147 c_comm = None
148 positional_zero = args and args[-1] == 0
149 if positional_zero or kwargs.get("dyn_comm") == 0 or kwargs.get("dyn") == 0:
150 c_comm = 0
151 dyn_dp_comm = kwargs.get("dyn_dp_comm", c_comm)
152 dyn_tp_comm = kwargs.get("dyn_tp_comm", c_comm)
153 dyn_cp_comm = kwargs.get("dyn_cp_comm", c_comm)
154 dyn_ep_comm = kwargs.get("dyn_ep_comm", c_comm)
155 dyn_ep_comm_balanced = kwargs.get("dyn_ep_comm_balanced", None)
156 dyn_ep_comm_imbalanced = kwargs.get("dyn_ep_comm_imbalanced", None)
157 if not self.__is_valid_eval_func(dyn_dp_comm):
158 dyn_dp_comm = self._ctx.node_eval[target_node].dyn.comm.dp
159 if not self.__is_valid_eval_func(dyn_tp_comm):
160 dyn_tp_comm = self._ctx.node_eval[target_node].dyn.comm.tp
161 if not self.__is_valid_eval_func(dyn_cp_comm):
162 dyn_cp_comm = self._ctx.node_eval[target_node].dyn.comm.cp
163 if not self.__is_valid_eval_func(dyn_ep_comm):
164 dyn_ep_comm = self._ctx.node_eval[target_node].dyn.comm.ep
165 comm_cls_obj = cls_obj
166 if self.is_regular_layer(target_node):
167 comm_cls_obj = EvalLayerComm
168 ep_balanced = (
169 self.__custom_getattr(comm_cls_obj, dyn_ep_comm_balanced)
170 if self.__is_valid_eval_func(dyn_ep_comm_balanced)
171 else None
172 )
173 ep_imbalanced = (
174 self.__custom_getattr(comm_cls_obj, dyn_ep_comm_imbalanced)
175 if self.__is_valid_eval_func(dyn_ep_comm_imbalanced)
176 else None
177 )
178 return NodeCommEval(
179 self.__custom_getattr(comm_cls_obj, dyn_dp_comm),
180 self.__custom_getattr(comm_cls_obj, dyn_tp_comm),
181 self.__custom_getattr(comm_cls_obj, dyn_cp_comm),
182 self.__custom_getattr(comm_cls_obj, dyn_ep_comm),
183 ep_balanced=ep_balanced,
184 ep_imbalanced=ep_imbalanced,
185 )
187 def __resolve_compute_fun(self, name):
188 """Resolve compute formula name — bypasses __wrap_mem_counter.
190 Compute FLOPs are not memory bytes and must not be accumulated
191 into accu_mem_type. This method directly retrieves the function
192 from EvalExpertCompute without the mem_counter wrapper.
193 Falls back to a zero-returning function if the name is not found,
194 so YAML typos produce zero FLOPs instead of crashing.
195 """
196 fun = getattr(EvalExpertCompute, name, None)
197 if fun is not None:
198 return fun
199 logger.warning("Compute formula '%s' not found in EvalExpertCompute, "
200 "falling back to zero", name)
202 def zero(*_):
203 return 0
205 zero.__qualname__ = f"0({name})"
206 return zero
208 def __set_node_eval_compute_fun(self, target_node, **kwargs):
209 """Build NodeComputeEval from kwargs, falling back to existing values."""
210 compute_kwargs = kwargs.get("compute", None)
211 if compute_kwargs is None:
212 # No compute config provided — preserve existing or return None
213 existing = self._ctx.node_eval.get(target_node)
214 if existing and existing.dyn.compute is not None:
215 return existing.dyn.compute
216 return None
217 if not isinstance(compute_kwargs, dict):
218 return None
219 router_name = compute_kwargs.get("router", None)
220 expert_balanced_name = compute_kwargs.get("expert_balanced", None)
221 expert_imbalanced_name = compute_kwargs.get("expert_imbalanced", None)
222 shared_expert_name = compute_kwargs.get("shared_expert", None)
223 # Resolve from EvalExpertCompute (bypass mem_counter)
224 router = self.__resolve_compute_fun(router_name) if router_name else None
225 expert_balanced = (
226 self.__resolve_compute_fun(expert_balanced_name)
227 if expert_balanced_name else None
228 )
229 expert_imbalanced = (
230 self.__resolve_compute_fun(expert_imbalanced_name)
231 if expert_imbalanced_name else None
232 )
233 shared_expert = (
234 self.__resolve_compute_fun(shared_expert_name)
235 if shared_expert_name else None
236 )
237 if not any([router, expert_balanced, expert_imbalanced, shared_expert]):
238 return None
239 return NodeComputeEval(
240 router=router,
241 expert_balanced=expert_balanced,
242 expert_imbalanced=expert_imbalanced,
243 shared_expert=shared_expert,
244 )
246 def set_head_eval_fun(self, *arg, **kwarg):
247 """overwrite head formulas"""
248 self.__set_node_eval_fun(EvalHead, self._ctx.head_node, *arg, **kwarg)
250 def set_tail_eval_fun(self, *arg, **kwarg):
251 """overwrite tail formulas"""
252 self.__set_node_eval_fun(EvalTail, self._ctx.tail_node, *arg, **kwarg)
254 def set_body_eval_fun(self, *args, **kwargs):
255 """overwrite body formulas"""
256 lay_type = kwargs.get("lay_type", args[0] if args else None)
257 if not lay_type:
258 lt = [
259 b_obj
260 for b_obj in list(LayerType)
261 if self.is_regular_layer(b_obj)
262 ]
263 else:
264 if not isinstance(lay_type, LayerType):
265 b_obj = self.__custom_getattr(LayerType, lay_type)
266 lt = [b_obj]
267 else:
268 lt = [lay_type]
269 for b_obj in lt:
270 self.__set_node_eval_fun(EvalBody, b_obj, *args, **kwargs)
272 def set_attn_eval_fun(
273 self,
274 num_p: Any = None,
275 qkv: Any = None,
276 score: Any = None,
277 proj: Any = None,
278 ) -> None:
279 """overwrite attention formulas"""
280 if self.__is_valid_eval_func(num_p):
281 self._ctx.attn_num_p = self.__custom_getattr(EvalAttn, num_p)
282 if self.__is_valid_eval_func(qkv):
283 self._ctx.attn_qkv_activ = self.__custom_getattr(
284 EvalAttn, qkv, MemType.ATTN_ACTIV
285 )
286 if self.__is_valid_eval_func(score):
287 self._ctx.attn_score_activ = self.__custom_getattr(
288 EvalAttn, score, MemType.ATTN_ACTIV
289 )
290 if self.__is_valid_eval_func(proj):
291 self._ctx.attn_proj_activ = self.__custom_getattr(
292 EvalAttn, proj, MemType.ATTN_ACTIV
293 )
295 def set_ffn_eval_fun(self, num_p: Any = None, activation=None, moe_activ=None):
296 """overwrite feedforward formulas"""
297 if self.__is_valid_eval_func(num_p):
298 self._ctx.ffn_num_p = self.__custom_getattr(EvalFFn, num_p)
299 if self.__is_valid_eval_func(activation):
300 self._ctx.ffn_activ = self.__custom_getattr(
301 EvalFFn, activation, MemType.FFN_ACTIV
302 )
303 if self.__is_valid_eval_func(moe_activ):
304 self._ctx.ffn_moe_activ = self.__custom_getattr(
305 EvalFFn, moe_activ, MemType.FFN_ACTIV
306 )
308 def set_expert_param_eval_fun(
309 self, routed_num_p: Any = None, shared_num_p: Any = None, router_num_p: Any = None
310 ) -> None:
311 """overwrite expert param count formulas for routed/shared breakdown, and the router's"""
312 if self.__is_valid_eval_func(routed_num_p):
313 self._ctx.ffn_routed_num_p = self.__custom_getattr(
314 EvalFFn, routed_num_p
315 )
316 if self.__is_valid_eval_func(shared_num_p):
317 self._ctx.ffn_shared_num_p = self.__custom_getattr(
318 EvalFFn, shared_num_p
319 )
320 if self.__is_valid_eval_func(router_num_p):
321 self._ctx.ffn_router_num_p = self.__custom_getattr(EvalFFn, router_num_p)
323 def set_norm_eval_fun(self, num_p: Any = None, activation=None):
324 """overwrite norm formulas"""
325 if self.__is_valid_eval_func(num_p):
326 self._ctx.norm_num_p = self.__custom_getattr(EvalNorm, num_p)
327 if self.__is_valid_eval_func(activation):
328 self._ctx.norm_activ = self.__custom_getattr(
329 EvalNorm, activation, MemType.NORM_ACTIV
330 )
332 def set_pp_micro_factor_eval_fun(self, sched_name, fun):
333 """overwrite PP microfactor formulas"""
334 if sched_name and self.__is_valid_eval_func(fun):
335 self._ctx.pp_micro_eval[sched_name] = self.__custom_getattr(
336 EvalUtils, fun
337 )
339 # Cost Model Config setter
341 def set_strategy(self, **kwargs):
342 """overwrite parallelism"""
343 self._ccfg.set_strategy(**kwargs)
344 self.fetch_hook_if_unimodal()
346 def set_ccfg(self, hook):
347 """overwrite cost model variable (except strategy)"""
348 if hook and callable(hook):
349 with strategy_guarded(self._ccfg):
350 hook(self._ccfg)
352 def __wrap_mem_counter(self, mem_type: MemType, fun: Callable) -> None:
353 """Wrap formula calls to accumulate memory by type."""
354 if mem_type and not hasattr(fun, "wrapped_with_counter"):
356 def wrap(*args, **kwargs):
357 res = fun(*args, **kwargs)
358 self._ctx.accu_mem_type[mem_type] += res
359 self._ctx.save2log(mem_type, res)
360 return res
362 wrap.__qualname__ = fun.__qualname__
363 wrap.wrapped_with_counter = True
364 return wrap
365 return fun
367 def __custom_getattr(
368 self, eval_class: Any, field: Any, mem_type: MemType = None
369 ) -> Callable:
370 """formula retrieve/wrap"""
371 # Definition priority order :
372 # 1. Callable (user defined in code)
373 # OR Numeric value
374 # 2. Overriding list from cost_model_preprocess (ccfg attribute)
375 # 3. Function name (config_eval yaml)
376 res = None
377 if callable(field):
378 res = field
379 if self.toggle_func_trace == field.__name__:
380 res = self.func_tracer.wrap(field)
381 if isinstance(field, (int, float)):
383 def constant(*_):
384 return field
386 constant.__qualname__ = str(field)
387 res = constant
388 if field in self._ccfg.overwrite_eval_functions:
389 res = self._ccfg.overwrite_eval_functions[field]
390 if isinstance(field, str):
392 def zero(*_):
393 return 0
395 zero.__qualname__ = "0"
396 res = getattr(eval_class, field, zero)
397 if self.toggle_func_trace == field:
398 res = self.func_tracer.wrap(getattr(eval_class, field))
399 res = self.__wrap_mem_counter(mem_type, res)
400 if not res:
401 raise TypeError(f"In eval config yaml, non valid field: {field}")
402 return res
404 # Import Eval Config
406 def import_eval_yaml(self) -> None:
407 """import evaluator config file, init ctx (inner call only)"""
408 if not self.toggle_func_trace:
409 self.toggle_func_trace = self.eval_cfg.trace_fun
410 # head
411 h_obj = getattr(LayerType, self.eval_cfg.nodes_mem_comp.head.name)
412 self._ctx.head_node = h_obj
413 self.set_head_eval_fun(
414 num_p=self.eval_cfg.nodes_mem_comp.head.num_param_fun,
415 stat_p=self.eval_cfg.nodes_mem_comp.head.stat_fun.p,
416 stat_os=self.eval_cfg.nodes_mem_comp.head.stat_fun.os,
417 stat_grad=self.eval_cfg.nodes_mem_comp.head.stat_fun.grad,
418 dyn_activ=self.eval_cfg.nodes_mem_comp.head.dyn_fun.activation,
419 dyn_dp_comm=self.eval_cfg.nodes_mem_comp.head.dyn_fun.comm.dp,
420 dyn_tp_comm=self.eval_cfg.nodes_mem_comp.head.dyn_fun.comm.tp,
421 dyn_cp_comm=self.eval_cfg.nodes_mem_comp.head.dyn_fun.comm.cp,
422 dyn_ep_comm=self.eval_cfg.nodes_mem_comp.head.dyn_fun.comm.ep,
423 )
425 # tail
426 t_obj = getattr(LayerType, self.eval_cfg.nodes_mem_comp.tail.name)
427 self._ctx.tail_node = t_obj
428 self.set_tail_eval_fun(
429 num_p=self.eval_cfg.nodes_mem_comp.tail.num_param_fun,
430 stat_p=self.eval_cfg.nodes_mem_comp.tail.stat_fun.p,
431 stat_os=self.eval_cfg.nodes_mem_comp.tail.stat_fun.os,
432 stat_grad=self.eval_cfg.nodes_mem_comp.tail.stat_fun.grad,
433 dyn_activ=self.eval_cfg.nodes_mem_comp.tail.dyn_fun.activation,
434 dyn_dp_comm=self.eval_cfg.nodes_mem_comp.tail.dyn_fun.comm.dp,
435 dyn_tp_comm=self.eval_cfg.nodes_mem_comp.tail.dyn_fun.comm.tp,
436 dyn_cp_comm=self.eval_cfg.nodes_mem_comp.tail.dyn_fun.comm.cp,
437 dyn_ep_comm=self.eval_cfg.nodes_mem_comp.tail.dyn_fun.comm.ep,
438 )
440 # body
441 for b in self.eval_cfg.nodes_mem_comp.body:
442 b_cfg = Config(b)
443 comm_cfg = b_cfg.dyn_fun.comm
444 comm_kwargs = {
445 "dyn_dp_comm": comm_cfg.dp,
446 "dyn_tp_comm": comm_cfg.tp,
447 "dyn_cp_comm": comm_cfg.cp,
448 "dyn_ep_comm": comm_cfg.ep,
449 }
450 if hasattr(comm_cfg, "ep_balanced"):
451 comm_kwargs["dyn_ep_comm_balanced"] = comm_cfg.ep_balanced
452 if hasattr(comm_cfg, "ep_imbalanced"):
453 comm_kwargs["dyn_ep_comm_imbalanced"] = comm_cfg.ep_imbalanced
454 # Compute config (FLOPs, not memory)
455 compute_kwargs = {}
456 compute_cfg = getattr(b_cfg.dyn_fun, "compute", None)
457 if compute_cfg and hasattr(compute_cfg, "__dict__"):
458 compute_kwargs = {"compute": dict(vars(compute_cfg))}
459 self.set_body_eval_fun(
460 lay_type=b_cfg.name,
461 num_p=b_cfg.num_param_fun,
462 stat_p=b_cfg.stat_fun.p,
463 stat_os=b_cfg.stat_fun.os,
464 stat_grad=b_cfg.stat_fun.grad,
465 dyn_activ=b_cfg.dyn_fun.activation,
466 **comm_kwargs,
467 **compute_kwargs,
468 )
470 # pp micro factor
471 for sc in self.eval_cfg.pp_sched:
472 self.set_pp_micro_factor_eval_fun(sc["name"], sc["fun"])
473 if not self._ccfg.pp_sched:
474 self._ccfg.pp_sched = self.eval_cfg.default_pp_sched
476 # layerblock
477 self.set_attn_eval_fun(
478 self.eval_cfg.base_arch_mem_comp.attention.num_param_fun,
479 self.eval_cfg.base_arch_mem_comp.attention.qkv,
480 self.eval_cfg.base_arch_mem_comp.attention.score,
481 self.eval_cfg.base_arch_mem_comp.attention.proj,
482 )
483 self.set_ffn_eval_fun(
484 self.eval_cfg.base_arch_mem_comp.feedforward.num_param_fun,
485 self.eval_cfg.base_arch_mem_comp.feedforward.activation,
486 self.eval_cfg.base_arch_mem_comp.feedforward.moe_activ,
487 )
488 self.set_expert_param_eval_fun(
489 routed_num_p=self.eval_cfg.base_arch_mem_comp.feedforward.routed_num_fun,
490 shared_num_p=self.eval_cfg.base_arch_mem_comp.feedforward.shared_num_fun,
491 router_num_p=self.eval_cfg.base_arch_mem_comp.feedforward.router_num_fun,
492 )
493 self.set_norm_eval_fun(
494 self.eval_cfg.base_arch_mem_comp.norm.num_param_fun,
495 self.eval_cfg.base_arch_mem_comp.norm.activation,
496 )
498 # passes
499 self.set_passes(
500 vpp_less_mem=self.eval_cfg.passes.vpp_less_memory,
501 swap_os=self.eval_cfg.passes.swap_optimizer,
502 dropless_tok_factor=self.eval_cfg.passes.dropless_tok_factor,
503 )
505 self._ctx.comm_expr = self.eval_cfg.comm_expr
507 def is_regular_layer(self, lay):
508 """check if layer is not head/tail"""
509 if isinstance(lay, str):
510 return lay[0] not in [
511 self._ctx.head_node.name[0],
512 self._ctx.tail_node.name[0],
513 ]
514 return lay not in [self._ctx.head_node, self._ctx.tail_node]