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

1# pylint: skip-file 

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

3 

4 

5class Template(MemEvalHook): 

6 

7 # Formula to get hooked 

8 @staticmethod 

9 def f(ccfg, ctx): 

10 return None 

11 

12 # Overwriting cost model variable 

13 @staticmethod 

14 def custom_ccfg(ccfg): 

15 pass 

16 

17 @staticmethod 

18 @hook_runner("model name") 

19 def run_hooks(e): 

20 c = Template 

21 e.set_ccfg(c.custom_ccfg) 

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

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

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

25 e.set_body_eval_fun( 

26 "NOT_REC_LAYER", 

27 num_p=c.f, 

28 stat_p=c.f, 

29 stat_os=c.f, 

30 stat_grad=c.f, 

31 dyn_activ=c.f, 

32 dyn_dp_comm=c.f, 

33 dyn_tp_comm=c.f, 

34 dyn_cp_comm=c.f, 

35 dyn_ep_comm=c.f, 

36 ) 

37 e.set_attn_eval_fun( 

38 num_p=c.f, 

39 qkv=c.f, 

40 score=c.f, 

41 proj=c.f, 

42 ) 

43 e.set_ffn_eval_fun( 

44 num_p=c.f, 

45 activation=c.f, 

46 moe_activ=c.f 

47 ) 

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

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