Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / memory_estimation / hooks / template.py: 100%

1# Copyright 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"""Skeleton of a memory estimation hook class. 

16 

17Copy the class, give @hook_runner the model name to hook, and replace the 

18placeholder formula and cost model override. 

19""" 

20from typing import Any 

21 

22from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.hook_base import MemEvalHook, hook_runner 

23 

24 

25class Template(MemEvalHook): 

26 """Registers a placeholder for every hookable formula.""" 

27 

28 @staticmethod 

29 def f(ccfg: Any, ctx: Any) -> None: 

30 """Formula to get hooked. The placeholder returns None.""" 

31 del ccfg, ctx 

32 

33 @staticmethod 

34 def custom_ccfg(ccfg: Any) -> None: 

35 """Cost model variables to overwrite. The placeholder changes nothing.""" 

36 del ccfg 

37 

38 @staticmethod 

39 @hook_runner("model name") 

40 def run_hooks(e: Any) -> None: 

41 """Register the placeholders on the evaluator e.""" 

42 c = Template 

43 e.set_ccfg(c.custom_ccfg) 

44 e.set_passes(vpp_less_mem=False, swap_os=False, dropless_tok_factor=1) 

45 e.set_head_eval_fun(num_p=c.f, stat=c.f, dyn=c.f) 

46 e.set_tail_eval_fun(num_p=c.f, stat=c.f, dyn=c.f) 

47 e.set_body_eval_fun( 

48 "NOT_REC_LAYER", 

49 num_p=c.f, 

50 stat_p=c.f, 

51 stat_os=c.f, 

52 stat_grad=c.f, 

53 dyn_activ=c.f, 

54 dyn_dp_comm=c.f, 

55 dyn_tp_comm=c.f, 

56 dyn_cp_comm=c.f, 

57 dyn_ep_comm=c.f, 

58 ) 

59 e.set_attn_eval_fun( 

60 num_p=c.f, 

61 qkv=c.f, 

62 score=c.f, 

63 proj=c.f, 

64 ) 

65 e.set_ffn_eval_fun( 

66 num_p=c.f, 

67 activation=c.f, 

68 moe_activ=c.f 

69 ) 

70 e.set_norm_eval_fun(num_p=c.f, activation=c.f) 

71 e.set_pp_micro_factor_eval_fun("1f1b", c.f)