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

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"""Walkthrough of the EvaluatorV2 memory estimation API, see README.md.""" 

16from typing import Any 

17 

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

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

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

21from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.logger import logger 

22 

23 

24def my_attn_num_param(ccfg: Any, ctx: Any) -> float: 

25 """Attention parameter count that overrides the default formula.""" 

26 del ctx 

27 return 10 * ccfg.h * ccfg.h 

28 

29 

30def custom(ccfg: Any) -> None: 

31 """Cost model variables that override the parsed configuration.""" 

32 ccfg.bytes_compute = 1 

33 ccfg.s = 1024 

34 ccfg.n_attMM = 5 

35 

36 

37def main() -> None: 

38 """Estimate, inspect and customize the memory of the bundled test cases.""" 

39 # Instantiate evaluator with a model configuration, 

40 # log_level=0 removes warning messages 

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

42 

43 # Check all defined node type 

44 logger.output("%s", list(LayerType)) 

45 

46 # Estimate peak memory (in Megabytes) 

47 peak_mem = e.estimate_peak(verbose=True) 

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

49 e.mem_fit(peak_mem) 

50 

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

52 logger.output("%s", e.static_mem_stage(1)) 

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

54 logger.output("%s", e.dynamic_mem_stage(1)) 

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

56 logger.output("%s", e.static_mem_layer(LayerType.FULL_REC_LAYER, 1)) 

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

58 logger.output("%s", e.dynamic_mem_layer(LayerType.FULL_REC_LAYER, 1)) 

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

60 logger.output("%s", e.logs_mem_stage(1)) 

61 # Fetch memory insights from each pipeline stage 

62 logger.output("%s", e.estimate_peak_insight()) 

63 # PPB Input 

64 logger.output("%s", e.estimate_layer_memory()) 

65 

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

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

68 

69 # Plot 

70 e.estimate_peak(plot=True) 

71 

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

73 

74 # Overwriting context function 

75 e.set_attn_eval_fun(num_p=my_attn_num_param) 

76 

77 # Overwriting a training feature 

78 e.set_passes(swap_os=True) 

79 

80 # Overwriting cost model variables 

81 e.set_ccfg(custom) 

82 

83 # Overwriting strategy 

84 logger.output("%s", e.get_strategy()) 

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

86 logger.output("%s", e.get_strategy()) 

87 

88 e.estimate_peak(verbose=True) 

89 # Inspect ccfg object (cost model variables) 

90 e.print_ccfg() 

91 # Inspect ctx object (evaluation variables and functions) 

92 e.print_ctx() 

93 

94 # Load a hook class 

95 # ... when declaring an Evaluator 

96 e = EvaluatorV2( 

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

98 ) 

99 # ... by using load_hook_cls() 

100 e.load_hook_cls(Template()) 

101 

102 

103if __name__ == "__main__": 

104 main()