Coverage for /home/jenkins/.local/lib/python3.10/site-packages/hyper_parallel/auto_parallel/sapp_nd/perf_estimation/comm

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

7# http://www.apache.org/licenses/LICENSE-2.0

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"""Experimental : Comm time"""

16from copy import deepcopy

17from hyper_parallel.auto_parallel.sapp_nd.nd.logger import perf_logger as logger

18import hyper_parallel.auto_parallel.sapp_nd.nd.common.hardware as Hard

19import hyper_parallel.auto_parallel.sapp_nd.nd.dimensions as Dim

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

21from hyper_parallel.auto_parallel.sapp_nd.nd.debug import PerfParts

22from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.comm import EvalLayerComm

23from hyper_parallel.auto_parallel.sapp_nd.memory_estimation._context import NodeEval, Context

24from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.head import EvalHead

25from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.tail import EvalTail

26from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.body import EvalBody

27from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.layer_block import (

28 EvalAttn,

29 EvalFFn,

30 EvalNorm,

31)

32from hyper_parallel.auto_parallel.sapp_nd.perf_estimation.utils_classes import NetworkLevel, PerformanceType

33from hyper_parallel.auto_parallel.sapp_nd.perf_estimation.getters import (

34 get_layer_custom_configs,

35 get_table_quantity,

36)

38COUNT_OPTIMIZER = False

41def fill_dp_table(cfg, tables):

42 """DP"""

43 table_dp = {}

44 table_dp["n_attMM"] = cfg.h * cfg.h / cfg.t

45 table_dp["n_ffMM"] = cfg.h * cfg.hff / cfg.t

46 table_dp["n_normOp"] = 2 * cfg.h / cfg.sp

48 if COUNT_OPTIMIZER:

49 table_dp["n_attParamCast"] = (

50 11 * cfg.h * cfg.h / (cfg.d if cfg.has_op else 1)

51 )

52 table_dp["n_ffParamCast"] = (

53 11 * cfg.h * cfg.hff / (cfg.d if cfg.has_op else 1)

54 )

55 for op in table_dp:

56 table_dp[op] *= cfg.bytes_norm if op == "n_normOp" else cfg.bytes_p

58 table_exp_dp = deepcopy(table_dp)

59 table_exp_dp["n_ffMM"] = (

60 2

61 * (cfg.n_exp + cfg.n_shared_exp)

62 * cfg.h

63 * cfg.hff_exp

64 / cfg.t

65 * cfg.bytes_p

66 )

67 tables[Dim.DP] = table_dp

68 tables["exp_dp"] = table_exp_dp

71def fill_tp_table(cfg, tables):

72 """TP"""

73 table_tp = {}

74 high_tp_bias = 11 / 16 if cfg.t >= 8 else 1 # Fix this

75 table_tp["n_gather"] = cfg.b * cfg.s * cfg.h * high_tp_bias

77 for op in table_tp:

78 table_tp[op] *= cfg.bytes_compute

80 table_exp_tp = deepcopy(table_tp)

81 table_exp_tp["n_gather"] = (

82 cfg.b * cfg.s * cfg.h * 1.5 * (cfg.ep / cfg.d) * cfg.bytes_compute

83 )

84 tables["tp"] = table_tp

85 tables["exp_tp"] = table_exp_tp

88def fill_ep_table(cfg, tables, device_type):

89 """EP"""

90 intra_devices = device_type.intra_node_num()

91 table_ep = {}

92 inter_node_bias_ep = 1

93 table_ep["n_ffMM"] = (

94 4

95 * cfg.n_chosen_exp

96 * cfg.b

97 * cfg.s

98 * cfg.h

99 * (max(4, cfg.os_max_shard) / cfg.t)

100 * cfg.cap_fact

101 * (

102 cfg.os_max_shard / min(intra_devices, cfg.ep)

103 + (

104 inter_node_bias_ep

105 * cfg.os_max_shard

106 / (cfg.ep / intra_devices)

107 if cfg.ep > intra_devices

108 else 0

109 )

110 )

111 )

112

113 for op in table_ep:

114 table_ep[op] *= cfg.bytes_compute

115 tables[Dim.EP] = table_ep

116

117

118def dp_ratio(cfg, device_type):

119 """formula"""

120 return (

121 0

122 if cfg.comm_d_non_exp == 0

123 else 1

124 - True # overlap_dp, Completely overlap standard DP comm

125 + (

126 1 / 16

127 if cfg.n_exp == 1

128 else 1 / max(1, cfg.ep / device_type.intra_node_num()) / 1.25

129 ) # overlap_op, Bias in overlapping OP comm (todo:make it dynamic too)

130 * (cfg.comm_d_non_exp - 1)

131 * cfg.os_max_shard

132 / cfg.d

133 )

134

135

136def comm_embed_ouput(cfg):

137 """ "formula"""

138 comm_embed = cfg.bytes_compute * cfg.h * cfg.v / cfg.shard_embed

139 comm_output = cfg.h * cfg.v / cfg.t

140 return comm_embed, comm_output

141

142

143def estimate_op_bulk_comm(*args, **kwargs):

144 """FW + BW"""

145 param = {

146 "cfg": args[0],

147 "ccfg": args[1],

148 "stages": args[2],

149 "device_type": args[3],

150 "with_recomp": kwargs.get(

151 "with_recomp", args[4] if len(args) > 4 else False

152 ),

153 "debugger": kwargs.get("debugger", args[5] if len(args) > 5 else None),

154 }

155

156 param["tables"] = {}

157 fill_dp_table(param["cfg"], param["tables"])

158

159 param['dp_ratio'] = dp_ratio(param['cfg'], param['device_type'])

160

161 param["comm_embed"], param["comm_output"] = comm_embed_ouput(param["cfg"])

162

163 if param["cfg"].dc_kv != 0: # Deepseek

164 param["comm_output"] += param["cfg"].h * (

165 2 * param["cfg"].h + param["cfg"].v

166 )

167 param["comm_output"] *= param["cfg"].n_mtp

168

169 param["comm_output"] *= param["cfg"].bytes_p

170

171 fill_tp_table(param["cfg"], param["tables"])

172 fill_ep_table(param["cfg"], param["tables"], param["device_type"])

173

174 lccfgs = get_layer_custom_configs(param["cfg"])

175 logger.info(lccfgs)

176 param["layer_count"] = 0

177 param["idx_lccfg"] = 0

178 comms = {Dim.DP: [], Dim.TP: [], Dim.EP: []}

179 # ignores comm recomp, to improve

180 for stage in param["stages"]:

181 comm = {Dim.DP: 0.0, Dim.TP: 0.0, Dim.EP: 0.0}

182 for chunk in stage:

183 for layer in chunk:

184 param["layer_count"], param["idx_lccfg"] = (

185 estimate_op_bulk_comm_layer(

186 param,

187 lccfgs,

188 layer=layer,

189 layer_count=param["layer_count"],

190 idx_lccfg=param["idx_lccfg"],

191 )

192 )

193 if param["ccfg"].ttype == PerformanceType.TIME:

194 for dim, ov in zip([Dim.DP, Dim.TP, Dim.DP], [0.0, 0.0, 0.0]):

195 comm[dim] = estimate_comm_score(

196 param["cfg"],

197 comm[dim],

198 dim,

199 overlap=ov,

200 device=param["device_type"],

201 )

202

203 comm[Dim.DP] *= param["dp_ratio"]

204 comm[Dim.TP] *= param["cfg"].comm_t

205 comm[Dim.EP] *= param["cfg"].comm_ep

206

207 if param["device_type"].name == "A3":

208 logger.info("A3 ratio")

209 comm[Dim.TP] /= 3

210

211 comms[Dim.DP].append(comm[Dim.DP])

212 comms[Dim.TP].append(comm[Dim.TP])

213 comms[Dim.EP].append(comm[Dim.EP])

214

215 if param["debugger"] and param["debugger"].is_enabled():

216 logger.info("DP_COMM = %s", comms[Dim.DP])

217 logger.info("MP_COMM = %s", comms[Dim.TP])

218 logger.info("EP_COMM = %s", comms[Dim.EP])

219 param["debugger"].info[PerfParts.DP_COMM] = comms[Dim.DP]

220 param["debugger"].info[PerfParts.MP_COMM] = comms[Dim.TP]

221 param["debugger"].info[PerfParts.EP_COMM] = comms[Dim.EP]

222

223 res = []

224 for i, c in enumerate(comms[Dim.TP]):

225 res.append(comms[Dim.DP][i] + c + comms[Dim.EP][i])

226

227 return res

228

229

230def estimate_op_bulk_comm_layer(cfg, lccfgs, **kwargs):

231 """for estimate_op_bulk_comm"""

232 if kwargs["layer"] == LayerType.EMBEDDING_LAYER:

233 kwargs["comm"][Dim.DP] += kwargs["param"]["comm_embed"]

234 return kwargs["layer_count"]

235

236 if kwargs["layer"] == LayerType.OUTPUT_LAYER:

237 kwargs["comm"][Dim.DP] += kwargs["param"]["comm_output"]

238 if cfg.dc_kv != 0: # Deepseek

239 lccfg = lccfgs[kwargs["idx_lccfg"]][0]

240 kwargs["comm"][Dim.TP] += cfg.n_mtp * get_table_quantity(

241 lccfg,

242 kwargs["param"]["tables"]["exp_tp"],

243 LayerType.NOT_REC_LAYER,

244 kwargs["param"]["with_recomp"],

245 )

246 return kwargs["layer_count"]

247

248 if (

249 kwargs["idx_lccfg"] + 1 < len(lccfgs)

250 and lccfgs[kwargs["idx_lccfg"]][1] == kwargs["layer_count"]

251 ):

252 kwargs["layer_count"] = 0

253 kwargs["idx_lccfg"] += 1

254

255 lccfg = lccfgs[kwargs["idx_lccfg"]][0]

256 is_moe_layer = lccfg.n_exp > 1

257

258 if is_moe_layer:

259 kwargs["comm"][Dim.DP] += get_table_quantity(

260 lccfg,

261 kwargs["param"]["tables"]["exp_dp"],

262 kwargs["layer"],

263 kwargs["param"]["with_recomp"],

264 )

265 kwargs["comm"][Dim.TP] += get_table_quantity(

266 lccfg,

267 kwargs["param"]["tables"]["exp_tp"],

268 kwargs["layer"],

269 kwargs["param"]["with_recomp"],

270 )

271 kwargs["comm"][Dim.EP] += get_table_quantity(

272 lccfg,

273 kwargs["param"]["tables"][Dim.EP],

274 kwargs["layer"],

275 kwargs["param"]["with_recomp"],

276 )

277 else:

278 kwargs["comm"][Dim.DP] += get_table_quantity(

279 lccfg,

280 kwargs["param"]["tables"][Dim.DP],

281 kwargs["layer"],

282 kwargs["param"]["with_recomp"],

283 )

284 kwargs["comm"][Dim.TP] += get_table_quantity(

285 lccfg,

286 kwargs["param"]["tables"]["tp"],

287 kwargs["layer"],

288 kwargs["param"]["with_recomp"],

289 )

290

291 kwargs["layer_count"] += 1

292 return kwargs["layer_count"], kwargs["idx_lccfg"]

293

294

295def prepare_context():

296 """context object"""

297 ctx = Context()

298 ctx.attn_num_p = EvalAttn.num_params_attn

299 ctx.ffn_num_p = EvalFFn.num_params_ffn

300 ctx.norm_num_p = EvalNorm.num_params_norm

301

302 ctx.node_eval[LayerType.EMBEDDING_LAYER] = NodeEval(

303 EvalHead.num_params_embed, None, None

304 )

305 ctx.node_eval[LayerType.OUTPUT_LAYER] = NodeEval(

306 EvalTail.num_params_output, None, None

307 )

308 ctx.node_eval[LayerType.NOT_REC_LAYER] = NodeEval(

309 EvalBody.num_params_layer, None, None

310 )

311 ctx.enable_accu_log = False

312 return ctx

313

314

315def estimate_from_mem_comm(*args, **kwargs):

316 """For memory estimation"""

317

318 param = {

319 "cfg": args[0],

320 "ccfg": args[1],

321 "stages": args[2],

322 "device_type": args[3],

323 }

324 param["debugger"] = kwargs.get(

325 "debugger", args[5] if len(args) > 5 else None

326 )

327 param["ctx"] = prepare_context()

328

329 # For layer type

330 param["flatten"] = sum(

331 [[f[1]] * f[0] for f in param["cfg"].layer_custom_config], []

332 )

333 comms = {Dim.DP: [], Dim.TP: [], Dim.EP: [], Dim.CP: []}

334 for stage in param["stages"]:

335 comm = {Dim.DP: 0.0, Dim.TP: 0.0, Dim.EP: 0.0, Dim.CP: 0.0}

336 for chunk in stage:

337 for layer in chunk:

338 param["ctx"].current_node = layer

339 if (

340 layer

341 not in [LayerType.EMBEDDING_LAYER, LayerType.OUTPUT_LAYER]

342 and param["flatten"]

343 ):

344 custom_fun = param["flatten"].pop(0)

345 if custom_fun:

346 custom_fun(param["cfg"])

347 logger.info("is layer moe ? %s", param["cfg"].n_exp > 1)

348 param["ctx"].current_node = LayerType.NOT_REC_LAYER

349 logger.info("param ctx %s", param["ctx"])

350 comm[Dim.DP] += EvalLayerComm.dp_comm_layer(param["cfg"], param["ctx"])

351

352 comm[Dim.TP] += EvalLayerComm.tp_comm_layer(

353 param["cfg"], param["ctx"], 1

354 ) # / 4 #* (param["cfg"].t - 1)

355 comm[Dim.EP] += EvalLayerComm.ep_comm_layer(

356 param["cfg"], param["ctx"], 1

357 ) # * param["cfg"].ep

358 comm[Dim.CP] += EvalLayerComm.cp_comm_layer(

359 param["cfg"], param["ctx"]

360 )

361 # min(device_type.level_bound_number[0], param["cfg"].ep)

362 # comm_cp += EvalLayerComm.cp_comm_layer

363 # (param["cfg"], param["ctx"])

364

365

366

367 if param["ccfg"].ttype == PerformanceType.TIME:

368 for dim, ov in zip([Dim.DP, Dim.TP, Dim.DP], [0.9, 0, 0.0]):

369 comm[dim] = estimate_comm_score(

370 param["cfg"],

371 comm[dim],

372 dim,

373 overlap=ov,

374 device=param["device_type"],

375 )

376

377 # if (

378 # not param["cfg"].gmm

379 # and param["cfg"].ep > param["device_type"].level_bound_number[0]

380 # ):

381 # comm[Dim.EP] *= 8

382

383 # if comm[Dim.EP] > 0:

384 # comm[Dim.EP] *= 2

385 # comm[Dim.TP] /= 2

386

387 # comm[Dim.DP] /= 100

388

389 # comm[Dim.EP] += comm[Dim.EP] * param["cfg"].ep / 100

390 # comm[Dim.TP] += comm[Dim.TP] * param["cfg"].t / 100

391

392 dev_per_node = param["device_type"].level_bound_number[0]

393 comm[Dim.TP] *= max(1, param["cfg"].t // dev_per_node)

394 comm[Dim.EP] *= max(1, param["cfg"].ep // dev_per_node)

395 comm[Dim.CP] *= max(1, param["cfg"].cp // dev_per_node)

396

397 comm[Dim.TP] /= 2 # TO REMOVE: FOR RUIWEN TEST ONLY

398 comm[Dim.DP] = 0 # /= 10 # TO REMOVE: FOR RUIWEN TEST ONLY

399

400 if param["device_type"].name == "A3":

401 logger.info("A3 ratio")

402 comm[Dim.DP] /= 2

403 comm[Dim.TP] /= 2

404 comm[Dim.EP] /= 2

405 comm[Dim.CP] /= 2

406

407 comms[Dim.DP].append(comm[Dim.DP])

408 comms[Dim.TP].append(comm[Dim.TP])

409 comms[Dim.EP].append(comm[Dim.EP])

410 comms[Dim.CP].append(comm[Dim.CP])

411

412 if param["debugger"] and param["debugger"].is_enabled():

413 logger.info("DP_COMM = %s", comms[Dim.DP])

414 logger.info("MP_COMM = %s", comms[Dim.TP])

415 logger.info("EP_COMM = %s", comms[Dim.EP])

416 logger.info("CP_COMM = %s", comms[Dim.CP])

417 param["debugger"].info[PerfParts.DP_COMM] = comms[Dim.DP]

418 param["debugger"].info[PerfParts.MP_COMM] = comms[Dim.TP]

419 param["debugger"].info[PerfParts.EP_COMM] = comms[Dim.EP]

420 param["debugger"].info[PerfParts.CP_COMM] = comms[Dim.CP]

421

422 res = []

423 for i, c in enumerate(comms[Dim.TP]):

424 res += [c + comms[Dim.DP][i] + comms[Dim.EP][i] + comms[Dim.CP][i]]

425

426 return res

427

428

429def estimate_comm(*args, **kwargs):

430 """wrapper"""

431 cfg, ccfg, stages, device_type = args[0], args[1], args[2], args[3]

432 with_recomp = kwargs.get(

433 "with_recomp", args[4] if len(args) > 4 else False

434 )

435 debugger = kwargs.get("debugger", args[5] if len(args) > 5 else None)

436 # return estimate_op_bulk_comm(cfg, ccfg, stages,

437 # device_type=device_type, with_recomp=with_recomp,

438 # debugger=debugger)

439 return estimate_from_mem_comm(

440 cfg,

441 ccfg,

442 stages,

443 device_type,

444 with_recomp=with_recomp,

445 debugger=debugger,

446 )

447

448

449def level_efficiency(level):

450 """to improve for Ascend A2"""

451 if level == NetworkLevel.NODE:

452 return 0.7

453 if level == NetworkLevel.CLUSTER:

454 return 0.9

455 raise ValueError

456

457

458def level_bandwidth(level):

459 """to improve for Ascend A2"""

460 if level == NetworkLevel.NODE:

461 return 300

462 if level == NetworkLevel.CLUSTER:

463 return 25

464 raise ValueError

465

466

467def level_latency(level):

468 """to improve for Ascend A2"""

469 if level == NetworkLevel.NODE:

470 return 0.00001

471 if level == NetworkLevel.CLUSTER:

472 return 0.00002

473 raise ValueError

474

475

476def comm_throughput(level):

477 """formula"""

478 eff = level_efficiency(level)

479 bw = level_bandwidth(level)

480 return bw * eff

481

482

483def estimate_comm_size_time(_, comm_size, level):

484 """formula"""

485 th = comm_throughput(level)

486 lat = level_latency(level)

487 return lat + comm_size / th

488

489

490def estimate_comm_score(

491 cfg, comm_volume, dim, overlap=0.0, device=Hard.device_map["A2"]

492):

493 """score assignment"""

494 assignment = device.level_assign(dp=cfg.d, tp=cfg.t, pp=cfg.p)

495 score = 0

496 for level in range(device.levels):

497 # intra_comm = comm_volume * (1-overlap)

498 # * (assignment[dim][0]-1) / device.intra_node_bw

499 score += (

500 comm_volume

501 * (1 - overlap)

502 * (

503 (assignment[dim][level] - 1)

504 * device.devices_below_level(level)

505 / device.level_bandwidth[level]

506 )

507 )

508 return score

Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / perf_estimation / comm_time.py: 87%

222 statements