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

1# pylint: skip-file 

2from hyper_parallel.auto_parallel.sapp_nd.nd.common.layer_type import LayerType 

3from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.estimate_v2 import EvaluatorV2 

4from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.hooks.template import Template 

5 

6# Instantiate evaluator with a model configuration, 

7# log_level=0 removes warning messages 

8e = EvaluatorV2("./test_cases/mixtral/default.yaml", log_level=0) 

9 

10# Check all defined node type 

11print(list(LayerType)) 

12 

13# Estimate peak memory (in Megabytes) 

14peak_mem = e.estimate_peak(verbose=True) 

15# Check whether estimation fits in device's max memory 

16e.mem_fit(peak_mem) 

17 

18# Estimate static memory of a specific pipeline stage (in Megabytes) 

19print(e.static_mem_stage(1)) 

20# Estimate dynamic memory of a specific pipeline stage (in Megabytes) 

21print(e.dynamic_mem_stage(1)) 

22# Estimate static memory of a specific layer and stage (in Megabytes) 

23print(e.static_mem_layer(LayerType.FULL_REC_LAYER, 1)) 

24# Estimate dynamic memory of a specific layer and stage (in Megabytes) 

25print(e.dynamic_mem_layer(LayerType.FULL_REC_LAYER, 1)) 

26# Retrieve the memory estimation logs of a specific stage (in Megabytes) 

27print(e.logs_mem_stage(1)) 

28# Fetch memory insights from each pipeline stage 

29stage_insights = e.estimate_peak_insight() 

30print(stage_insights) 

31# PPB Input 

32ppb_input = e.estimate_layer_memory() 

33print(ppb_input) 

34 

35# Inspect a specific stage (here is the first one) 

36e.estimate_peak(spec_stage_id=0, verbose=True) 

37 

38# Plot 

39e.estimate_peak(plot=True) 

40 

41e = EvaluatorV2("./test_cases/deepseek3/default.yaml", log_level=0) 

42 

43 

44# Overwriting context function 

45def my_attn_num_param(ccfg, ctx): 

46 return 10 * ccfg.h * ccfg.h 

47 

48 

49e.set_attn_eval_fun(num_p=my_attn_num_param) 

50 

51# Overwriting a training feature 

52e.set_passes(swap_os=True) 

53 

54 

55# Overwriting cost model variables 

56def custom(ccfg): 

57 ccfg.bytes_compute = 1 

58 ccfg.s = 1024 

59 ccfg.n_attMM = 5 

60 

61 

62e.set_ccfg(custom) 

63 

64# Overwriting strategy 

65print(e.get_strategy()) 

66e.set_strategy(dp=8, tp=8, m=128) 

67print(e.get_strategy()) 

68 

69e.estimate_peak(verbose=True) 

70# Inspect ccfg object (cost model variables) 

71e.print_ccfg() 

72# Inspect ctx object (evaluation variables and functions) 

73e.print_ctx() 

74 

75# Load a hook class 

76# ... when declaring an Evaluator 

77e = EvaluatorV2( 

78 "./test_cases/deepseek3/default.yaml", log_level=0, hook_cls=Template() 

79) 

80# ... by using load_hook_cls() 

81e.load_hook_cls(Template())