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

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"""Context module for evaluator""" 

16from __future__ import annotations 

17from pprint import pformat 

18from typing import TYPE_CHECKING 

19from dataclasses import dataclass 

20from enum import Enum, auto 

21from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.utils import EvalUtils 

22 

23if TYPE_CHECKING: 

24 from typing import Self, Any 

25 

26 

27class MemType(Enum): 

28 """memory types""" 

29 

30 MODEL_PARAM = auto() 

31 OPTIM_STATE = auto() 

32 ACCU_GRAD = auto() 

33 ATTN_ACTIV = auto() 

34 FFN_ACTIV = auto() 

35 NORM_ACTIV = auto() 

36 AG_COMM = auto() 

37 A2A_COMM = auto() 

38 

39 

40@dataclass 

41class NodeStatEval: 

42 """static formula pointers""" 

43 

44 p: Any 

45 os: Any 

46 grad: Any 

47 

48 def __repr__(self): 

49 return ( 

50 f"stat.p={_qname(self.p)}, " 

51 f"stat.os={_qname(self.os)}, " 

52 f"stat.grad={_qname(self.grad)}" 

53 ) 

54 

55 

56def _qname(attr): 

57 """Safe qualname accessor for __repr__ — handles None and non-callable values.""" 

58 return getattr(attr, "__qualname__", str(attr)) 

59 

60 

61@dataclass 

62class NodeCommEval: 

63 """comm formula pointers""" 

64 

65 dp: Any 

66 tp: Any 

67 cp: Any 

68 ep: Any 

69 ep_balanced: Any = None 

70 ep_imbalanced: Any = None 

71 

72 def __repr__(self): 

73 return ( 

74 f"dyn.comm.dp={_qname(self.dp)}, " 

75 f"dyn.comm.tp={_qname(self.tp)}, " 

76 f"dyn.comm.cp={_qname(self.cp)}, " 

77 f"dyn.comm.ep={_qname(self.ep)}" 

78 ) 

79 

80 

81@dataclass 

82class NodeDynEval: 

83 """dynamic formula pointers""" 

84 

85 activation: Any 

86 comm: NodeCommEval 

87 

88 def __repr__(self): 

89 return f"dyn.activation={_qname(self.activation)}, " f"{str(self.comm)}" 

90 

91 

92@dataclass 

93class NodeEval: 

94 """Associate a LayerType -> 

95 (Num param function, static mem function, dynamic mem function) 

96 """ 

97 

98 num_p: Any 

99 stat: NodeStatEval 

100 dyn: NodeDynEval 

101 

102 def __repr__(self): 

103 return ( 

104 f"num_p = {self.num_p.__name__}, " 

105 f"{str(self.stat)}, " 

106 f"{str(self.dyn)}" 

107 ) 

108 

109 

110class Context: 

111 """Context class""" 

112 

113 def __init__(self) -> None: 

114 """initializing buffers""" 

115 # Temporary bufferes 

116 self.enable_node_log = True 

117 self.accu_mem_type = {mt: 0 for mt in list(MemType)} 

118 self.node_compute_log = {} 

119 

120 # Map node to (static function, dynamic function) 

121 self.node_eval = {} 

122 # Variables 

123 self.vpp_less_mem, self.swap_os = None, None 

124 self.dropless_tok_factor = None 

125 self.attn_num_p, self.attn_qkv_activ = None, None 

126 self.attn_score_activ, self.attn_proj_activ = None, None 

127 self.ffn_num_p, self.ffn_activ, self.ffn_moe_activ = None, None, None 

128 self.ffn_routed_num_p, self.ffn_shared_num_p = None, None 

129 self.norm_num_p, self.norm_activ = None, None 

130 self.pp_micro_eval = {} 

131 self.head_node, self.tail_node = None, None 

132 self.current_node = None 

133 self.current_stage_id, self.current_chunk_id = -1, -1 

134 self.current_lay_id = None 

135 self.real_lay_ids = [] 

136 self.ppb, self.default_micro_factor = None, None 

137 

138 def __str__(self): 

139 return pformat( 

140 dict( 

141 (k, v) if k != "node_eval" else (k, self.print_node_eval()) 

142 for k, v in vars(self).items() 

143 ) 

144 ) 

145 

146 def print_node_eval(self): 

147 """from all layertype""" 

148 return dict((k, str(v)) for k, v in self.node_eval.items()) 

149 

150 @property 

151 def eval(self): 

152 """shortcut""" 

153 return self.node_eval[self.current_node] 

154 

155 def init_tmp_buff(self) -> None: 

156 """reset""" 

157 self.enable_node_log = True 

158 self.accu_mem_type = {mt: 0 for mt in list(MemType)} 

159 self.node_compute_log = {} 

160 

161 def copy_tmp_buff(self, target_ctx: Self) -> None: 

162 """copy to target_ctx""" 

163 for att, val in vars(self).items(): 

164 if att != "node_eval": 

165 setattr(target_ctx, att, val) 

166 

167 def save2log(self, fun, val_in_bytes): 

168 """lay_id -> fun, val""" 

169 if self.enable_node_log and val_in_bytes > 0: 

170 name = fun 

171 if callable(fun): 

172 name = fun.__name__ 

173 elif isinstance(fun, MemType): 

174 name = fun.name.lower() 

175 node_name = self.current_node 

176 if not isinstance(self.current_node, str): 

177 node_name = self.current_node.name[0] 

178 if isinstance(self.current_lay_id, int): 

179 real_lay_id = self.real_lay_ids[self.current_chunk_id][ 

180 self.current_stage_id 

181 ][self.current_lay_id] 

182 else: 

183 lay_id = int(self.current_lay_id.split("_")[-1]) 

184 real_lay_id = self.real_lay_ids[self.current_chunk_id][ 

185 self.current_stage_id 

186 ][lay_id] 

187 real_lay_id = self.current_lay_id.replace( 

188 str(lay_id), str(real_lay_id) 

189 ) 

190 # Add key 

191 pair = ( 

192 self.current_stage_id, 

193 self.current_chunk_id, 

194 real_lay_id, 

195 node_name, 

196 ) 

197 if pair not in self.node_compute_log: 

198 self.node_compute_log[pair] = {} 

199 if name not in self.node_compute_log[pair]: 

200 self.node_compute_log[pair][name] = 0 

201 self.node_compute_log[pair][name] += EvalUtils.mb(val_in_bytes)