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 NodeComputeEval: 

83 """compute formula pointers""" 

84 

85 router: Any = None 

86 expert_balanced: Any = None 

87 expert_imbalanced: Any = None 

88 shared_expert: Any = None 

89 

90 def __repr__(self): 

91 parts = [] 

92 if self.router is not None: 

93 parts.append(f"compute.router={_qname(self.router)}") 

94 if self.expert_balanced is not None: 

95 parts.append(f"compute.expert_balanced={_qname(self.expert_balanced)}") 

96 if self.expert_imbalanced is not None: 

97 parts.append(f"compute.expert_imbalanced={_qname(self.expert_imbalanced)}") 

98 if self.shared_expert is not None: 

99 parts.append(f"compute.shared_expert={_qname(self.shared_expert)}") 

100 return ", ".join(parts) if parts else "compute=None" 

101 

102 

103@dataclass 

104class NodeDynEval: 

105 """dynamic formula pointers""" 

106 

107 activation: Any 

108 comm: NodeCommEval 

109 compute: NodeComputeEval = None 

110 

111 def __repr__(self): 

112 s = f"dyn.activation={_qname(self.activation)}, {str(self.comm)}" 

113 if self.compute is not None: 

114 s += f", {str(self.compute)}" 

115 return s 

116 

117 

118@dataclass 

119class NodeEval: 

120 """Associate a LayerType -> 

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

122 """ 

123 

124 num_p: Any 

125 stat: NodeStatEval 

126 dyn: NodeDynEval 

127 

128 def __repr__(self): 

129 return ( 

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

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

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

133 ) 

134 

135 

136class Context: 

137 """Context class""" 

138 

139 def __init__(self) -> None: 

140 """initializing buffers""" 

141 # Temporary bufferes 

142 self.enable_node_log = True 

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

144 self.node_compute_log = {} 

145 

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

147 self.node_eval = {} 

148 # Variables 

149 self.vpp_less_mem, self.swap_os = None, None 

150 self.dropless_tok_factor = None 

151 self.attn_num_p, self.attn_qkv_activ = None, None 

152 self.attn_score_activ, self.attn_proj_activ = None, None 

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

154 self.ffn_routed_num_p, self.ffn_shared_num_p = None, None 

155 self.norm_num_p, self.norm_activ = None, None 

156 self.pp_micro_eval = {} 

157 self.head_node, self.tail_node = None, None 

158 self.current_node = None 

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

160 self.current_lay_id = None 

161 self.real_lay_ids = [] 

162 self.ppb, self.default_micro_factor = None, None 

163 

164 def __str__(self): 

165 return pformat( 

166 dict( 

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

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

169 ) 

170 ) 

171 

172 def print_node_eval(self): 

173 """from all layertype""" 

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

175 

176 @property 

177 def eval(self): 

178 """shortcut""" 

179 return self.node_eval[self.current_node] 

180 

181 def init_tmp_buff(self) -> None: 

182 """reset""" 

183 self.enable_node_log = True 

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

185 self.node_compute_log = {} 

186 

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

188 """copy to target_ctx""" 

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

190 if att != "node_eval": 

191 setattr(target_ctx, att, val) 

192 

193 def save2log(self, fun, val_in_bytes): 

194 """lay_id -> fun, val""" 

195 if self.enable_node_log and val_in_bytes > 0: 

196 name = fun 

197 if callable(fun): 

198 name = fun.__name__ 

199 elif isinstance(fun, MemType): 

200 name = fun.name.lower() 

201 node_name = self.current_node 

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

203 node_name = self.current_node.name[0] 

204 if isinstance(self.current_lay_id, int): 

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

206 self.current_stage_id 

207 ][self.current_lay_id] 

208 else: 

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

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

211 self.current_stage_id 

212 ][lay_id] 

213 real_lay_id = self.current_lay_id.replace( 

214 str(lay_id), str(real_lay_id) 

215 ) 

216 # Add key 

217 pair = ( 

218 self.current_stage_id, 

219 self.current_chunk_id, 

220 real_lay_id, 

221 node_name, 

222 ) 

223 if pair not in self.node_compute_log: 

224 self.node_compute_log[pair] = {} 

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

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

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