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

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"""PPB input module""" 

16 

17from hyper_parallel.auto_parallel.sapp_nd.nd.common.layer_type import LayerType 

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

19 

20 

21class _PPB: 

22 """Pipeline balance payload builder.""" 

23 

24 def __init__(self, eval_cfg, inner_dyn_fun): 

25 self.eval_cfg = eval_cfg 

26 self._inner_dynamic_mem = inner_dyn_fun 

27 self.mb = EvalUtils.mb 

28 

29 def add_to_ppb_list(self, ppb_lay_desc, desc): 

30 """layer description list preparation""" 

31 if desc: 

32 already_comp = False 

33 body_idx = 0 

34 for d in ppb_lay_desc: 

35 if all(v == d[k] for k, v in desc.items()): 

36 # already exist desc 

37 d["nb_layer"] += 1 

38 already_comp = True 

39 if d["type"] == "BODY": 

40 body_idx += 1 

41 if desc and not already_comp: 

42 desc["nb_layer"] = 1 

43 if desc["type"] == "BODY": 

44 desc["name"] = f"BODY_{body_idx}" 

45 else: 

46 desc["name"] = desc["type"] 

47 ppb_lay_desc += [desc] 

48 

49 def lay_ppb(self, ccfg, ctx, res_stat): 

50 """layer description preparation""" 

51 ctx.enable_node_log = False 

52 desc = {} 

53 desc["model_name"] = ccfg.model_name 

54 if ctx.current_node == ctx.head_node: 

55 d_emb = self.mb(sum(self._inner_dynamic_mem(ppb=True))) 

56 desc["type"] = "HEAD" 

57 desc["memory_parameter"] = self.mb(res_stat) + d_emb 

58 # desc["time"] = perfs[0] 

59 desc["time"] = 1 

60 elif ctx.current_node == ctx.tail_node: 

61 d_out = self.mb(sum(self._inner_dynamic_mem(ppb=True))) 

62 desc["type"] = "TAIL" 

63 desc["memory_parameter"] = self.mb(res_stat) + d_out 

64 # desc["time"] = perfs[-1] 

65 desc["time"] = 1 

66 else: 

67 ctx.current_node = LayerType.NOT_REC_LAYER 

68 dyn_nrec = self._inner_dynamic_mem(ppb=True) 

69 ctx.current_node = LayerType.SEL_REC_LAYER 

70 dyn_srec = self._inner_dynamic_mem(ppb=True) 

71 ctx.current_node = LayerType.FULL_REC_LAYER 

72 dyn_frec = self._inner_dynamic_mem(ppb=True) 

73 c = max(dyn_nrec[1], dyn_srec[1], dyn_frec[1]) # Temporary 

74 desc["type"] = "BODY" 

75 desc["memory_parameter"] = self.mb(res_stat) 

76 desc["memory_parameter"] += self.mb(c) 

77 desc["memory_activation"] = self.mb(dyn_nrec[0]) 

78 desc["memory_select_rec"] = self.mb(dyn_srec[0]) 

79 desc["memory_recompute"] = self.mb(dyn_frec[0]) 

80 desc["time"] = 1 

81 ctx.enable_node_log = True 

82 # if any( 

83 # k in desc and desc[k] > 0 

84 # for k in [ 

85 # "memory_parameter", 

86 # "memory_activation", 

87 # "memory_recompute", 

88 # "memory_select_rec", 

89 # ] 

90 # ): 

91 # # Ignore values 0 

92 # return desc 

93 return desc 

94 

95 def ppb_combine_bodies(self, ppb_lay_desc): 

96 """combine descriptions into a new body""" 

97 if not self.eval_cfg.ppb_combined: 

98 return 

99 for new_body in self.eval_cfg.ppb_combined: 

100 desc = { 

101 "model_name": "combined", 

102 "type": "BODY", 

103 "memory_parameter": 0, 

104 "memory_activation": 0, 

105 "memory_select_rec": 0, 

106 "memory_recompute": 0, 

107 "time": 1, 

108 "nb_layer": 1, 

109 "name": "COMBINED", 

110 } 

111 idx = -1 

112 for mod, t in new_body: 

113 target = next( 

114 ( 

115 d 

116 for d in ppb_lay_desc 

117 if d["model_name"] == mod and d["type"] == t.upper() 

118 ), 

119 None, 

120 ) 

121 if target: 

122 desc["model_name"] += "_" + mod 

123 desc["name"] += "_" + target["name"] 

124 for m in desc: 

125 if m.startswith("memory") and m in target: 

126 desc[m] += target[m] 

127 target_idx = ppb_lay_desc.index(target) 

128 idx = target_idx if idx < 0 else min(idx, target_idx) 

129 del ppb_lay_desc[target_idx] 

130 idx = max(idx, 0) 

131 ppb_lay_desc.insert(idx, desc) 

132 

133 def lay_ppb_new(self, ccfg, ctx, res_stat): 

134 """layer description preparation""" 

135 ctx.enable_node_log = False 

136 desc = {} 

137 desc["model_name"] = ccfg.model_name 

138 desc["memory_activation"] = {"NONE": 0, "FULL": 0} 

139 if ctx.current_node == ctx.head_node: 

140 d_emb = self.mb(sum(self._inner_dynamic_mem(ppb=True))) 

141 desc["memory_parameter"] = self.mb(res_stat) + d_emb 

142 desc["type"] = "HEAD" 

143 elif ctx.current_node == ctx.tail_node: 

144 d_out = self.mb(sum(self._inner_dynamic_mem(ppb=True))) 

145 desc["memory_parameter"] = self.mb(res_stat) + d_out 

146 desc["type"] = "TAIL" 

147 else: 

148 ctx.current_node = LayerType.NOT_REC_LAYER 

149 dyn_nrec = self._inner_dynamic_mem(ppb=True) 

150 ctx.current_node = LayerType.FULL_REC_LAYER 

151 dyn_frec = self._inner_dynamic_mem(ppb=True) 

152 c = max(dyn_nrec[1], dyn_frec[1]) # Temporary 

153 desc["memory_parameter"] = self.mb(res_stat) 

154 desc["memory_parameter"] += self.mb(c) 

155 desc["memory_activation"]["NONE"] = self.mb(dyn_nrec[0]) 

156 desc["memory_activation"]["FULL"] = self.mb(dyn_frec[0]) 

157 desc["type"] = "BODY" 

158 desc["options"] = ["NONE", "FULL"] 

159 desc["forward_time"] = {"NONE": 1, "FULL": 1} 

160 desc["backward_time"] = {"NONE": 1, "FULL": 1} 

161 desc["time"] = 1 

162 ctx.enable_node_log = True 

163 # if any( 

164 # k > 0 

165 # for k in [ 

166 # desc["memory_parameter"], 

167 # desc["memory_activation"]["NONE"], 

168 # desc["memory_activation"]["FULL"], 

169 # ] 

170 # ): 

171 # # Ignore values 0 

172 # return desc 

173 # return {} 

174 return desc 

175 

176 def ppb_combine_bodies_new(self, ppb_lay_desc): 

177 """combine descriptions into a new body""" 

178 if not self.eval_cfg.ppb_combined: 

179 return 

180 for new_body in self.eval_cfg.ppb_combined: 

181 desc = { 

182 "model_name": "combined", 

183 "type": "BODY", 

184 "memory_parameter": 0, 

185 "memory_activation": {"NONE": 0, "FULL": 0}, 

186 "options": ["NONE", "FULL"], 

187 "forward_time": {"NONE": 1, "FULL": 1}, 

188 "backward_time": {"NONE": 1, "FULL": 1}, 

189 "time": 1, 

190 "nb_layer": 1, 

191 "name": "COMBINED", 

192 } 

193 idx = -1 

194 for mod, t in new_body: 

195 target = next( 

196 ( 

197 d 

198 for d in ppb_lay_desc 

199 if d["model_name"] == mod and d["type"] == t.upper() 

200 ), 

201 None, 

202 ) 

203 if target: 

204 desc["model_name"] += "_" + mod 

205 desc["name"] += "_" + target["name"] 

206 desc["memory_parameter"] += target["memory_parameter"] 

207 desc["memory_activation"]["NONE"] += target[ 

208 "memory_activation" 

209 ]["NONE"] 

210 desc["memory_activation"]["FULL"] += target[ 

211 "memory_activation" 

212 ]["FULL"] 

213 target_idx = ppb_lay_desc.index(target) 

214 idx = target_idx if idx < 0 else min(idx, target_idx) 

215 del ppb_lay_desc[target_idx] 

216 idx = max(idx, 0) 

217 ppb_lay_desc.insert(idx, desc)