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

1# Copyright 2025 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"""getters, setters, printers""" 

16from __future__ import annotations 

17from typing import TYPE_CHECKING 

18 

19from hyper_parallel.auto_parallel.sapp_nd.memory_estimation._backbone import _Backbone 

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

21 

22if TYPE_CHECKING: 

23 from typing import Dict, Union, Tuple 

24 

25 

26class _Utils(_Backbone): 

27 """utils class""" 

28 

29 # def __init__(self, *args, **kwargs): 

30 # super().__init__(*args, **kwargs) 

31 

32 def get_model_name(self) -> str: 

33 """accessor""" 

34 return self._ccfg.model_name 

35 

36 def get_strategy(self) -> Dict: 

37 """return parallelism/recompute strategies""" 

38 return self._ccfg.get_strategy() 

39 

40 def get_max_device_memory(self) -> float: 

41 """accessor for max device memory in MB""" 

42 return self._ccfg.device_capacity.to_mb().size 

43 

44 def get_num_layers(self) -> Union[Tuple, int]: 

45 """tuple of all L if multimodal""" 

46 if not self._ccfg.multimodal: 

47 return self._ccfg.n_lay + self._ccfg.n_mtp 

48 return [ 

49 self._ccfg.mm_ccfgs[m].n_lay + self._ccfg.mm_ccfgs[m].n_mtp 

50 for m in self._ccfg.mm_order 

51 ] 

52 

53 def set_layer_custom(self, lc=None) -> None: 

54 """setting ccfg.layer_custom_config (inner call only)""" 

55 if not lc: 

56 self._ccfg.layer_custom_config = [(self._ccfg.n_lay, None)] 

57 elif isinstance(lc, list): 

58 self._ccfg.layer_custom_config = lc 

59 

60 def set_config(self, config) -> None: 

61 """Explicitly Assign a new config ccfg""" 

62 self._ccfg = config 

63 

64 def all_stage_micro_factors(self) -> None: 

65 """get all stage's warmup micros for current schedule""" 

66 sched = self._ccfg.pp_sched 

67 for stage_id in range(self._ccfg.p): 

68 chunk_micro = [] 

69 for chunk_id in range(self._ccfg.vp): 

70 self._ctx.current_stage_id = stage_id 

71 self._ctx.current_chunk_id = chunk_id 

72 micro = self._ctx.pp_micro_eval[sched](self._ccfg, self._ctx) 

73 chunk_micro += [micro] 

74 logger.info( 

75 "%s stage _%s = %s", 

76 self._ctx.pp_micro_eval[sched].__name__, 

77 stage_id, 

78 str(chunk_micro), 

79 ) 

80 

81 # Printers 

82 

83 def print_ccfg(self) -> None: 

84 """pretty printer for ccfg""" 

85 if not self._ccfg.multimodal: 

86 print(self._ccfg) 

87 else: 

88 for m in self._ccfg.mm_order: 

89 print("Module", m) 

90 print(self._ccfg.mm_ccfgs[m]) 

91 

92 def print_ctx(self) -> None: 

93 """pretty printer for ctx""" 

94 print("Eval Context attributes:\n" + str(self._ctx)) 

95 

96 def print_node_eval(self) -> None: 

97 """get all (P,stat,dyn)""" 

98 return self._ctx.print_node_eval() 

99 

100 def print_stages(self, stages: list, spec_stage_id: int = -1) -> None: 

101 """can target a stage id""" 

102 self._ccfg.print_stages(stages, spec_stage_id)