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.ffn_router_num_p = None 

156 self.norm_num_p, self.norm_activ = None, None 

157 self.pp_micro_eval = {} 

158 # How many chunks the schedule splits a micro-batch's sequence into, 

159 # as the micro factor of the node being evaluated counts them. 

160 self.seq_chunks = 1 

161 # How many layers' gathered parameters the node evaluated holds as 

162 # its working set in its backward, 0 when it is what it keeps 

163 # between its passes. 

164 self.working_set = 0 

165 # Whether the stage's dynamic memory counts what the node keeps for 

166 # the backward whose working set it is. 

167 self.working_on_saved = False 

168 # The recompute switches of the layer evaluated, where they are not 

169 # its config's rec_op: a layer priced under other switches. None: 

170 # the config's. 

171 self.switches = None 

172 # The gradient bytes of the node evaluated last. 

173 self.node_grad = 0 

174 self.head_node, self.tail_node = None, None 

175 self.current_node = None 

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

177 self.current_lay_id = None 

178 self.real_lay_ids = [] 

179 self.ppb, self.default_micro_factor = None, None 

180 

181 def __str__(self): 

182 return pformat( 

183 dict( 

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

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

186 ) 

187 ) 

188 

189 def print_node_eval(self): 

190 """from all layertype""" 

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

192 

193 @property 

194 def eval(self): 

195 """shortcut""" 

196 return self.node_eval[self.current_node] 

197 

198 def init_tmp_buff(self) -> None: 

199 """reset""" 

200 self.enable_node_log = True 

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

202 self.node_compute_log = {} 

203 

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

205 """copy to target_ctx""" 

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

207 if att != "node_eval": 

208 setattr(target_ctx, att, val) 

209 

210 def save2log(self, fun, val_in_bytes): 

211 """lay_id -> fun, val""" 

212 if self.enable_node_log and val_in_bytes > 0: 

213 name = fun 

214 if callable(fun): 

215 name = fun.__name__ 

216 elif isinstance(fun, MemType): 

217 name = fun.name.lower() 

218 node_name = self.current_node 

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

220 node_name = self.current_node.name[0] 

221 if isinstance(self.current_lay_id, int): 

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

223 self.current_stage_id 

224 ][self.current_lay_id] 

225 else: 

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

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

228 self.current_stage_id 

229 ][lay_id] 

230 real_lay_id = self.current_lay_id.replace( 

231 str(lay_id), str(real_lay_id) 

232 ) 

233 # Add key 

234 pair = ( 

235 self.current_stage_id, 

236 self.current_chunk_id, 

237 real_lay_id, 

238 node_name, 

239 ) 

240 if pair not in self.node_compute_log: 

241 self.node_compute_log[pair] = {} 

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

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

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