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

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_config(self, config) -> None: 

54 """Explicitly Assign a new config ccfg""" 

55 self._ccfg = config 

56 

57 def all_stage_micro_factors(self) -> None: 

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

59 sched = self._ccfg.pp_sched 

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

61 chunk_micro = [] 

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

63 self._ctx.current_stage_id = stage_id 

64 self._ctx.current_chunk_id = chunk_id 

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

66 chunk_micro += [micro] 

67 logger.info( 

68 "%s stage _%s = %s", 

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

70 stage_id, 

71 str(chunk_micro), 

72 ) 

73 

74 # Printers 

75 

76 def print_ccfg(self) -> None: 

77 """pretty printer for ccfg""" 

78 if not self._ccfg.multimodal: 

79 print(self._ccfg) 

80 else: 

81 for m in self._ccfg.mm_order: 

82 print("Module", m) 

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

84 

85 def print_ctx(self) -> None: 

86 """pretty printer for ctx""" 

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

88 

89 def print_node_eval(self) -> None: 

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

91 return self._ctx.print_node_eval() 

92 

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

94 """can target a stage id""" 

95 self._ccfg.print_stages(stages, spec_stage_id)