Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / nd / dimensions.py: 90%

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"""parallel dimensions""" 

16from __future__ import annotations 

17 

18import sys 

19from typing import TYPE_CHECKING 

20 

21from hyper_parallel.auto_parallel.sapp_nd.nd.logger import logger 

22 

23if TYPE_CHECKING: 

24 from hyper_parallel.auto_parallel.sapp_nd.nd.common._cost_model_variables import _CostModVar 

25 

26 

27class Dimension: 

28 """Output dimension""" 

29 

30 def __init__( 

31 self, 

32 acronym, 

33 cost_model_var_name, 

34 from_str, 

35 default=1, 

36 ): 

37 self.name = acronym 

38 self.cost_model_var = cost_model_var_name 

39 self.default = default 

40 self.bound = None 

41 self.from_str = from_str 

42 

43 def __str__(self): 

44 return self.name 

45 

46 def __repr__(self): 

47 return str(self) 

48 

49 def lname(self): 

50 """lower case name""" 

51 return self.name.lower() 

52 

53 def from_config(self, ccfg: _CostModVar): 

54 """Get dimension value from cost model config""" 

55 try: 

56 value = ccfg.__dict__[self.cost_model_var] 

57 except KeyError: 

58 logger.error( 

59 "variable %s does not exist in the cost model: %s", 

60 self.cost_model_var, 

61 str(ccfg.__dict__), 

62 ) 

63 sys.exit(1) 

64 return value 

65 

66 def reset_bound(self): 

67 """Reset bound the dimension space""" 

68 self.bound = None 

69 

70 def set_bound(self, bound): 

71 """Bound the dimension space""" 

72 if self.bound: 

73 logger.debug( 

74 "bound(%s) = min (%d, %d)", self.name, bound, self.bound 

75 ) 

76 self.bound = min(bound, self.bound) 

77 else: 

78 logger.debug("bound(%s) = %d", self.name, bound) 

79 self.bound = bound 

80 

81 def get_bound(self): 

82 """Return dimension bound""" 

83 return self.bound 

84 

85 def is_valid(self, value): 

86 """Check dimension value validity""" 

87 invalid = False 

88 if isinstance(value, bool): 

89 return not invalid 

90 if isinstance(value, int): 

91 invalid = self.bound and value > self.bound 

92 invalid = invalid or value < 1 

93 if invalid: 

94 logger.warning( 

95 "Dimension %s = %s is invalid", self.name, str(value) 

96 ) 

97 return not invalid 

98 

99 

100DP = Dimension( 

101 "DP", 

102 "d", 

103 default=1, 

104 from_str=int, 

105) 

106EP = Dimension( 

107 "EP", 

108 "ep", 

109 default=1, 

110 from_str=int, 

111) 

112TP = Dimension( 

113 "MP", 

114 "t", 

115 default=1, 

116 from_str=int, 

117) 

118CP = Dimension( 

119 "CP", 

120 "cp", 

121 default=1, 

122 from_str=int, 

123) 

124PP = Dimension( 

125 "PP", 

126 "p", 

127 default=1, 

128 from_str=int, 

129) 

130MBN = Dimension( 

131 "MB", 

132 "m", 

133 default=1, 

134 from_str=int, 

135) 

136MBS = Dimension( 

137 "MBS", 

138 "b", 

139 default=1, 

140 from_str=int, 

141) 

142SP = Dimension( 

143 "SP", 

144 "sp", 

145 default=True, 

146 from_str=bool, 

147) 

148OP = Dimension( 

149 "OP", 

150 "os_max_shard", 

151 # "op_weight_shard", 

152 default=1, 

153 from_str=int, 

154) 

155VPP = Dimension( 

156 "VPP", 

157 "vp", 

158 default=1, 

159 from_str=int, 

160) 

161 

162ALL_DIMS = [DP, EP, TP, CP, PP, VPP, MBN, MBS, SP, OP] 

163 

164 

165class Dimensions: 

166 """All output dimensions""" 

167 

168 def __init__(self, config, all_dims=None): 

169 if isinstance(config, list): 

170 self.all_dims = [d for d, _ in config] 

171 self.dims_val = dict(config) 

172 # elif isinstance(config, dict): 

173 # self.all_dims = ALL_DIMS 

174 # self.dims_val = {d: d.from_config(config) for d in self.all_dims} 

175 elif isinstance(config, bool): 

176 self.all_dims = ALL_DIMS 

177 self.dims_val = {d: d.default for d in self.all_dims} 

178 else: 

179 raise TypeError( 

180 f"Dimensions cannot be constructed from type {type(config)}" 

181 ) 

182 if all_dims: 

183 self.all_dims = all_dims 

184 self._reset_all_dims() 

185 

186 def _reset_all_dims(self): 

187 for d in self.all_dims: 

188 d.reset_bound() 

189 

190 def __str__(self): 

191 return str(self.dims_val) 

192 

193 def __repr__(self): 

194 return str(self) 

195 

196 def keys(self): 

197 """Return dimensions""" 

198 return list(self.dims_val) 

199 

200 def global_batch_size(self): 

201 """Compute the global batch size""" 

202 gbs = self.dims_val[DP] * self.dims_val[MBS] 

203 has_pp = self.has_dim(PP) and self.dims_val[PP] > 1 

204 has_mbn = self.has_dim(MBN) and self.dims_val[MBN] > 1 

205 if has_pp and has_mbn: 

206 gbs *= self.dims_val[MBN] 

207 return gbs 

208 

209 def values(self): 

210 """Return dimension value""" 

211 return [str(self.dims_val[d]) for d in self.dims_val] 

212 

213 def unique_name(self): 

214 """Return all values as a unique string""" 

215 return "_".join(self.values()) 

216 

217 def has_dim(self, d): 

218 """Check that this dimension has a value in the parallel config""" 

219 return d in self.dims_val 

220 

221 def is_valid(self): 

222 """Check if all dimensions values are valid""" 

223 if MBN in self.dims_val and PP in self.all_dims: 

224 valid = self.dims_val[MBN] >= self.dims_val[PP] 

225 valid = valid and not ( 

226 self.dims_val[PP] == 1 and self.dims_val[MBN] > 1 

227 ) 

228 if not valid: 

229 logger.warning("PP and MBN were deemed not suitable") 

230 return False 

231 if TP in self.all_dims and not ( 

232 self.dims_val[TP] & (self.dims_val[TP] - 1) == 0 

233 ): 

234 logger.warning("%s must be a power of 2", str(TP)) 

235 return False 

236 for d in self.dims_val: 

237 if not d.is_valid(self.dims_val[d]): 

238 logger.warning("Dimension %d is not valid", d) 

239 return False 

240 if SP in self.all_dims and CP in self.all_dims: 

241 if self.dims_val[SP] and self.dims_val[CP] > 1: 

242 logger.warning("SP & CP cannot coexist") 

243 return False 

244 if OP in self.all_dims: 

245 op = self.dims_val[OP] 

246 if not op & (op - 1) == 0: 

247 logger.warning("OP %d must be a power of 2", op) 

248 return False 

249 return True 

250 

251 def val(self, dim): 

252 """Get Dimension value""" 

253 return self.dims_val[dim] 

254 

255 def set(self, dim, val): 

256 """Get Dimension value""" 

257 self.dims_val[dim] = val 

258 

259 def steal(self, factor, dim_from, dim_to): 

260 """Assign a dimension factor to another dimension""" 

261 self.dims_val[dim_from] = self.dims_val[dim_from] // factor 

262 self.dims_val[dim_to] = self.dims_val[dim_to] * factor 

263 

264 

265def get_dim(acronym): 

266 """Return the dimension of the given string acronym""" 

267 dname = str(acronym).upper() 

268 for d in ALL_DIMS: 

269 if d.name == dname: 

270 return d 

271 raise ValueError(f"Dimension {dname} does NOT exist") 

272 

273 

274def get_dims(dims): 

275 """Return all dimensions considered""" 

276 if dims is None: 

277 return ALL_DIMS 

278 return [get_dim(acronym) for acronym in dims]