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

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"""performance estimation"""

16import json

17from copy import deepcopy

18import numpy as np

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

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

22from hyper_parallel.auto_parallel.sapp_nd.nd.common.cost_model_config import CostModelConfig

23from hyper_parallel.auto_parallel.sapp_nd.nd.common.arch_hooks import check_and_apply_custom_hook

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

25from hyper_parallel.auto_parallel.sapp_nd.nd.debug import PerfParts, RealParts, estimation_in_real_parts

27from hyper_parallel.auto_parallel.sapp_nd.perf_estimation.utils_classes import (

28 RatioType,

29 PerformanceType,

30 P2PCommType,

31 RecType,

32 CustomConfig,

33)

34from hyper_parallel.auto_parallel.sapp_nd.perf_estimation.comm_time import estimate_comm

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

36 get_layer_custom_configs,

37 get_table_quantity,

38)

41GENERALIZE_PIPELINE_CALCULATION = False

42MANUAL_P2P_RATIO = 0.002

43BACKWARD_RATIO = 2

46def op_table(cfg):

47 """op compute load formulas"""

48 table = {}

49 # cfg.s /= cfg.cp

50 table["n_attMM"] = (

51 3 * (1 + cfg.n_kv / cfg.a) * cfg.b * cfg.s * cfg.h * cfg.h

52 )

53 table["n_ffMM"] = 6 * cfg.b * cfg.s * cfg.h * cfg.hff

54 table["n_attBMM"] = 6 * cfg.b * cfg.s * cfg.s * cfg.h

55 table["n_ffBMM"] = 6 * cfg.b * cfg.s * cfg.s * cfg.hff

56 table["n_softmax"] = 13 * cfg.a * cfg.b * cfg.s * cfg.s

57 table["n_headCast"] = 3 * cfg.a * cfg.b * cfg.s * cfg.s

58 table["n_gather"] = cfg.b * cfg.s * cfg.h * (cfg.t - 1)

59 table["n_ffAct"] = 21 * cfg.b * cfg.hff

61 table["n_normOp"] = 30 * cfg.b * cfg.s * cfg.h * cfg.t / cfg.sp

62 table["n_dropout"] = (

63 3 * cfg.b * cfg.s * max(cfg.a * cfg.s, 3 * cfg.h * cfg.t / cfg.sp)

64 )

65 if cfg.dc_kv != 0: # Deepseek

66 table["n_attMM"] = (

67 3

68 / 2

69 * (

70 2 * cfg.dc_kv * cfg.n_kv * cfg.dh

71 + cfg.dc_q * cfg.a * (cfg.dh + cfg.dhr)

72 + cfg.h * (cfg.a * cfg.dh + cfg.dhr)

73 )

74 * cfg.b

75 * cfg.s

76 )

77 for op in table:

78 table[op] *= cfg.bytes_p / cfg.t / cfg.cp

79 # cfg.s *= cfg.cp

80 return table

83# Evaluation functions

84def estimate_op_bulk_comp(cfg, ccfg, stages, with_recomp=False, debugger=None):

85 """FW + BW"""

86 _ = debugger

87 table = op_table(cfg)

89 table_exp = deepcopy(table) # Verify this with MF MoEV2

90 table_exp["n_ffMM"] *= (

91 cfg.hff_exp / cfg.hff * max(1, cfg.n_chosen_exp) * cfg.cap_fact

92 )

93 table_exp["n_ffBMM"] *= (

94 cfg.hff_exp / cfg.hff * max(1, cfg.n_chosen_exp) * cfg.cap_fact

95 )

97 lccfgs = get_layer_custom_configs(cfg)

98 layer_count = 0

99 idx_lccfg = 0

100

101 flops = []

102 for stage in stages:

103 flops += [0]

104 for chunk in stage:

105 for layer in chunk:

106 if layer == LayerType.EMBEDDING_LAYER:

107 continue

108

109 if layer == LayerType.OUTPUT_LAYER:

110 flops[-1] += (1 if cfg.dc_kv == 0 else cfg.n_mtp) * (

111 1

112 / 16 # bias_imbalance

113 * 6

114 * cfg.b

115 * cfg.v

116 * cfg.h

117 * cfg.s

118 * cfg.bytes_p

119 / cfg.t

120 )

121 continue

122

123 layer_count += 1

124 if (

125 idx_lccfg + 1 < len(lccfgs)

126 and lccfgs[idx_lccfg][1] <= layer_count

127 ):

128 layer_count = 0

129 idx_lccfg += 1

130

131 flop = get_table_quantity(

132 lccfgs[idx_lccfg][0],

133 table_exp if (lccfgs[idx_lccfg][0].n_exp > 1) else table,

134 layer,

135 with_recomp,

136 )

137

138 if ccfg.ttype == PerformanceType.TIME:

139 flop = estimate_comp_flop_time(lccfgs[idx_lccfg][0], flop)

140

141 flops[-1] += flop

142

143 return flops

144

145

146def estimate_comp(cfg, ccfg, stages, with_recomp=False, debugger=None):

147 """wrapper"""

148 return estimate_op_bulk_comp(

149 cfg, ccfg, stages, with_recomp, debugger=debugger

150 )

151

152

153# Experimental : Flop time

154

155

156def efficiency(x):

157 """obtained via extrapolation"""

158 eff = min(

159 1.0, max(0.1, 0.00004694 * x**3 + 0.0014 * x**2 - 0.0336 * x + 0.1)

160 )

161 return eff

162

163

164def throughput(precision_bytes, flop):

165 """assumes matrix"""

166 eff = efficiency(flop / (10.0**12))

167 return precision_bytes**2 * (10.0**12) * eff

168

169

170def estimate_comp_flop_time(cfg, flop, is_softmax=False):

171 """flop from throughput"""

172 th = throughput(

173 cfg.bytes_softmax if is_softmax else cfg.bytes_compute, flop

174 )

175 return flop / th

176

177

178##############################################################

179

180

181def get_dynamic_ratio(cfg):

182 """comm/comp"""

183 if cfg.n_exp == 1:

184 return 3 / 2 * (cfg.hff + cfg.s) * (8192 / (cfg.h + cfg.s))

185 return 3 / 2 * (cfg.hff_exp + cfg.s) * (8192 / (cfg.h + cfg.s))

186

187

188def estimate_stage(*args, **kwargs):

189 """stage level estimation"""

190 cfg = args[0]

191 ccfg = args[1]

192 compute_perfs = args[2]

193 comm_perfs = args[3]

194 recompute_perfs = args[4]

195 recomm_perfs = args[5]

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

197 comp_w = 1

198 comm_w = 1

199 if ccfg.rtype == RatioType.COMM_ONLY:

200 comp_w = 0

201 elif ccfg.rtype == RatioType.COMPUTE_ONLY:

202 comm_w = 0

203 elif ccfg.rtype == RatioType.STATIC:

204 comm_w = 10**4

205 ccfg.static_ratio = comm_w

206 elif ccfg.rtype == RatioType.DYNAMIC:

207 comm_w = get_dynamic_ratio(cfg)

208 ccfg.dynamic_ratio = comm_w

209 perf = [

210 comp_w * compute_perfs[i] + comm_w * comm_perfs[i]

211 for i in range(len(compute_perfs))

212 ]

213 logger.info("ratio = %s", comm_w)

214 # ignores comm recomp, to improve

215 re_perf = [

216 (

217 max(0, comp_w * (recompute_perfs[i] - compute_perfs[i]))

218 + max(0, comm_w * (recomm_perfs[i] - comm_perfs[i]))

219 )

220 / (1 + BACKWARD_RATIO)

221 for i in range(len(compute_perfs))

222 ]

223

224 if debugger and debugger.is_enabled():

225 for p in [PerfParts.DP_COMM, PerfParts.MP_COMM, PerfParts.EP_COMM, PerfParts.CP_COMM]:

226 debugger.info[p] = [

227 comm_w * c for c in debugger.info[p]

228 ]

229 debugger.info[PerfParts.FW_COMPUTE] = [

230 comp_w * comp / (1 + BACKWARD_RATIO) for comp in compute_perfs

231 ]

232 debugger.info[PerfParts.BW_COMPUTE] = [

233 fw * BACKWARD_RATIO for fw in debugger.info[PerfParts.FW_COMPUTE]

234 ]

235 debugger.info[PerfParts.RECOMPUTE] = re_perf

236

237 return [perf[i] + re_perf[i] for i in range(len(perf))]

238 #penalty_fn(stage)

239 #return stage

240

241

242def estimate_pipeline(cfg, stage_perfs, stage_focused=None, debugger=None):

243 """pipeline level estimation"""

244 logger.info("stage_perfs = %s", stage_perfs)

245 straggler_time = max(stage_perfs)

246 sum_time = sum(stage_perfs)

247 last_straggler_idx = cfg.p - 1 - np.argmax(stage_perfs[::-1])

248 logger.info(

249 "straggler estim is %s and its stage is %s",

250 straggler_time,

251 last_straggler_idx,

252 )

253

254 non_steady_perf = 0

255 steady_perf = 0

256 if cfg.p == 1:

257 if len(stage_perfs) != 1:

258 raise ValueError("Expected exactly one stage performance")

259 steady_perf = sum_time * cfg.m

260 elif cfg.vp == 1:

261 non_steady_perf = sum_time

262 if GENERALIZE_PIPELINE_CALCULATION:

263 last_idx = last_straggler_idx + 1

264 steady_perf = (

265 cfg.m - cfg.p + last_straggler_idx

266 ) * straggler_time + sum(stage_perfs[last_idx:])

267 else:

268 steady_perf = (cfg.m - 1) * straggler_time

269 else:

270 less_extra = cfg.p * (cfg.vp - 1)

271 # big_extra = less_extra + cfg.p

272

273 # we assume that times of all micro-blocks in one vp chunk are the same

274 straggler_time /= cfg.vp

275 sum_time /= cfg.vp

276

277 # more_memory has more micro blocks in warm-up but it does

278 # not matter since they will overlap with steady phase

279 non_steady_perf = sum_time + less_extra * straggler_time

280

281 # more_memory is a boost to performance and a nerf to memory

282 # if cfg.vp_less_memory or True:

283 # steady_perf = (cfg.m * cfg.vp - less_extra - 1) * straggler_time

284 # else:

285 # steady_perf = (cfg.m * cfg.vp - big_extra - 1) * straggler_time

286 steady_perf = (cfg.m * cfg.vp - less_extra - 1) * straggler_time

287 straggler_time *= cfg.vp

288 sum_time *= cfg.vp

289

290 pipeline_perf = non_steady_perf + steady_perf

291 logger.info(

292 "pipeline_perf = non_steady_perf(%.2E) + steady_perf(%.2E)",

293 non_steady_perf,

294 steady_perf,

295 )

296

297 if stage_focused is not None:

298 last_straggler_idx = stage_focused

299 if debugger and debugger.is_enabled():

300 time_sum = 0

301 for k in [

302 PerfParts.DP_COMM,

303 PerfParts.MP_COMM,

304 PerfParts.EP_COMM,

305 PerfParts.CP_COMM,

306 PerfParts.FW_COMPUTE,

307 PerfParts.BW_COMPUTE,

308 PerfParts.RECOMPUTE,

309 ]:

310 debugger.info[k] = (

311 debugger.info[k][last_straggler_idx] * cfg.m

312 ) # / cfg.vp

313 logger.info(

314 "time_sum += debugger[%s] = %.2E",

315 k,

316 debugger.info[k],

317 )

318 time_sum += debugger.info[k]

319

320 if abs(time_sum - straggler_time * cfg.m) < 1e-9:

321 logger.warning("Inconsistency found in straggler time calculation")

322 time_sum = straggler_time * cfg.m

323 logger.info(

324 "straggler time = %.2E. %s x stragglers = %.2E",

325 straggler_time,

326 cfg.m,

327 straggler_time * cfg.m,

328 )

329

330 bubble = pipeline_perf - time_sum

331 logger.info(

332 "bubble(%.2E) = pipeline_perf(%.2E) - time_sum(%.2E)",

333 bubble,

334 pipeline_perf,

335 time_sum,

336 )

337 debugger.info[PerfParts.BUBBLE] = bubble

338 return pipeline_perf

339

340

341def estimate_p2p_comm(cfg, straggler, ratio=MANUAL_P2P_RATIO, debugger=None):

342 """pipeline comm"""

343 nb_send_recv = 0

344 if cfg.vp == 1:

345 nb_send_recv = (

346 0

347 if cfg.p == 1

348 else (

349 4 * cfg.m

350 if cfg.p == 2

351 else 4 * cfg.p * cfg.m + 4 * cfg.p * cfg.p - 14 * cfg.p

352 )

353 )

354 else:

355 nb_send_recv = (

356 0

357 if cfg.p == 1

358 else (

359 8 * cfg.m * cfg.vp - 4 * cfg.m

360 if cfg.p == 2

361 else (

362 16 * cfg.m * cfg.vp + 12

363 if cfg.p == 4

364 else 4 * cfg.p * cfg.m * cfg.vp

365 + 4 * cfg.p * cfg.p

366 - 13 * cfg.p

367 )

368 )

369 )

370 pp_comm = ratio * nb_send_recv / cfg.p * straggler / cfg.sp

371 if debugger and debugger.is_enabled():

372 debugger.info[PerfParts.PP_COMM] = pp_comm

373

374 return pp_comm

375

376

377def estimate_perf(cfg, _, stage_perfs, stage_focused=None, debugger=None):

378 """wrapper"""

379 return estimate_pipeline(cfg, stage_perfs, stage_focused=stage_focused, debugger=debugger)

380

381

382def estimate_p2p(cfg, ccfg, stage_perfs, debugger=None):

383 """wrapper"""

384 if ccfg.ptype != P2PCommType.MANUAL:

385 p2p = 0

386 else:

387 p2p = estimate_p2p_comm(cfg, max(stage_perfs), debugger=debugger)

388 if debugger and debugger.is_enabled():

389 debugger.info[PerfParts.PP_COMM] = p2p

390 return p2p

391

392

393def estimate_layer_perf(*args, **kwargs):

394 """for PPB"""

395 cfg = args[0]

396 stages = kwargs.get("stages", args[2] if len(args) > 2 else None)

397 extra_custom_func = kwargs.get(

398 "extra_custom_func", args[3] if len(args) > 3 else None

399 )

400 ccfg = kwargs.get("ccfg", args[4] if len(args) > 4 else CustomConfig())

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

402 # cfg = CostModelConfig(mf_config)

403 # Process custom model config

404 if extra_custom_func:

405 extra_custom_func(cfg)

406 else:

407 logger.info("auto applying custom model config")

408 check_and_apply_custom_hook(cfg)

409

410 new_layer_config = []

411 stages = [[LayerType.EMBEDDING_LAYER]]

412 for _, layer in cfg.layer_custom_config:

413 new_layer_config.append((1, layer))

414 stages.append([LayerType.NOT_REC_LAYER])

415 stages.append([LayerType.OUTPUT_LAYER])

416 cfg.layer_custom_config = new_layer_config

417

418 logger.output("cfg.layer_custom_config = %s", cfg.layer_custom_config)

419 logger.output("stages = %s", stages)

420

421 cfg.n = cfg.d * cfg.t * cfg.p

422

423 cfg.n_headCast = 1

424 cfg.n_ffAct = 1

425

426 logger.info(str(cfg))

427 logger.info(stages)

428 logger.info(ccfg)

429

430 perfs = {}

431 perfs["compute_perfs"] = estimate_comp(

432 cfg, ccfg, stages, with_recomp=False, debugger=debugger

433 )

434 logger.info("PerfEst: compute_perfs %s", perfs["compute_perfs"])

435 perfs["recompute_perfs"] = (

436 [0] * cfg.p

437 if ccfg.retype not in {RecType.COMPUTE_ONLY, RecType.WITH}

438 else estimate_comp(

439 cfg, ccfg, stages, with_recomp=True, debugger=debugger

440 )

441 )

442

443 perfs["comm_perfs"] = estimate_comm(

444 cfg, ccfg, stages, args[1], with_recomp=False, debugger=debugger

445 )

446 logger.info("PerfEst: comm_perfs %s", perfs["comm_perfs"])

447 perfs["recomm_perfs"] = (

448 [0] * cfg.p

449 if ccfg.retype not in {RecType.COMM_ONLY, RecType.WITH}

450 else estimate_comm(

451 cfg, ccfg, stages, args[1], with_recomp=True, debugger=debugger

452 )

453 )

454 stage_perfs = estimate_stage(

455 cfg,

456 ccfg,

457 perfs["compute_perfs"],

458 perfs["comm_perfs"],

459 perfs["recompute_perfs"],

460 perfs["recomm_perfs"],

461 debugger=debugger,

462 )

463 logger.output("PerfEst: stage_perfs %s", stage_perfs)

464

465 for s, perf in enumerate(stage_perfs):

466 stage_perfs[s] = int(perf / 10**12)

467

468 return stage_perfs

469

470

471

472def apply_regression_coefficients(coeffs, debugger, old_perf):

473 """

474 applies the coefficients present in regression's cache_file

475 """

476 compute_ratio = coeffs.get("COMPUTE")

477 for part, raw in list(debugger.info.items()):

478 if part in (PerfParts.TOTAL, PerfParts.MEMORY):

479 continue

480 if part in (PerfParts.FW_COMPUTE,

481 PerfParts.BW_COMPUTE,

482 PerfParts.RECOMPUTE):

483 ratio = compute_ratio

484 else:

485 ratio = coeffs.get(part.name)

486 new_val = 0.0 if raw == 0.0 else raw * ratio

487 debugger.info[part] = new_val

488

489 max_idx = max(p.value for p in PerfParts) -1

490 estimations = [0.0] * max_idx

491 for part in PerfParts:

492 if part in (PerfParts.TOTAL, PerfParts.MEMORY):

493 continue

494 estimations[part.value - 1] = debugger.info.get(part) or 0.0

495

496 real_buckets = {rp: [] for rp in RealParts}

497 real_buckets = estimation_in_real_parts(real_buckets, estimations, old_perf)

498

499 perf = (

500 real_buckets[RealParts.COMP][-1]

501 + real_buckets[RealParts.DP_WAIT][-1]

502 + real_buckets[RealParts.MP_WAIT][-1]

503 + real_buckets[RealParts.EP_WAIT][-1]

504 + real_buckets[RealParts.CP_WAIT][-1]

505 + real_buckets[RealParts.PP_WAIT][-1]

506 )

507 debugger.info[PerfParts.TOTAL] = perf

508 return perf

509

510

511def _resolve_estimate_args(args, kwargs):

512 """Resolve positional/keyword inputs for estimate_performance.

513

514 Returns:

515 Tuple of (cfg, stages, extra_custom_func, ccfg, debugger,

516 device_type, memory).

517 """

518 cfg_input = args[0]

519 cfg = (

520 cfg_input

521 if isinstance(cfg_input, CostModelConfig)

522 else CostModelConfig(cfg_input)

523 )

524 stages = kwargs.get("stages", args[1] if len(args) > 1 else None)

525 extra_custom_func = kwargs.get(

526 "extra_custom_func", args[2] if len(args) > 2 else None

527 )

528 ccfg = kwargs.get("ccfg", args[3] if len(args) > 3 else CustomConfig())

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

530 device_type = kwargs.get(

531 "device_type", args[5] if len(args) > 5 else Hard.device_map["A2"]

532 )

533 memory = kwargs.get("memory", args[6] if len(args) > 6 else None)

534 return cfg, stages, extra_custom_func, ccfg, debugger, device_type, memory

535

536

537def _finalize_perf(perf, cache_file, debugger, memory):

538 """Apply cached regression coefficients and record debug info.

539

540 Returns:

541 The final performance value.

542 """

543 coeffs = None

544 cache = False

545 if cache_file is not None:

546 with open(cache_file, 'r', encoding='utf-8') as f:

547 coeffs = json.load(f)

548 cache = True

549

550 if debugger and debugger.is_enabled():

551 if cache:

552 perf = apply_regression_coefficients(coeffs, debugger, perf)

553 debugger.info[PerfParts.TOTAL] = perf

554 if memory is not None:

555 debugger.info[PerfParts.MEMORY] = memory

556 return perf

557

558

559# performance estimation

560def estimate_performance(*args, **kwargs):

561 """main estimation"""

562 (

563 cfg,

564 stages,

565 extra_custom_func,

566 ccfg,

567 debugger,

568 device_type,

569 memory,

570 ) = _resolve_estimate_args(args, kwargs)

571

572 # Process custom model config

573 if extra_custom_func:

574 extra_custom_func(cfg)

575 else:

576 logger.info("auto applying custom model config")

577 check_and_apply_custom_hook(cfg)

578

579 # Process partition generation

580 if not stages:

581 logger.info("stage partitions are generated")

582 stages = cfg.generate_partitions_vpp()

583

584 cfg.n = cfg.d * cfg.t * cfg.p

585 cfg.n_headCast = 1

586 cfg.n_ffAct = 1

587

588 logger.debug(

589 "perf_model: DP = %d, TP = %d, EP = %d, PP = %d, MB = %d",

590 cfg.d,

591 cfg.t,

592 cfg.ep,

593 cfg.p,

594 cfg.m,

595 )

596

597 logger.info(str(cfg))

598 logger.info(stages)

599 logger.info(ccfg)

600

601 compute_perfs = estimate_comp(

602 cfg, ccfg, stages, with_recomp=False, debugger=debugger

603 )

604 recompute_perfs = (

605 [0] * cfg.p

606 if ccfg.retype not in {RecType.COMPUTE_ONLY, RecType.WITH}

607 else estimate_comp(

608 cfg, ccfg, stages, with_recomp=True, debugger=debugger

609 )

610 )

611 comm_perfs = estimate_comm(

612 cfg, ccfg, stages, device_type, with_recomp=False, debugger=debugger

613 )

614 logger.info("PerfEst: comm_perfs %s", comm_perfs)

615 recomm_perfs = (

616 [0] * cfg.p

617 if ccfg.retype not in {RecType.COMM_ONLY, RecType.WITH}

618 else estimate_comm(

619 cfg, ccfg, stages, device_type, with_recomp=True, debugger=debugger

620 )

621 )

622

623 stage_perfs = estimate_stage(

624 cfg,

625 ccfg,

626 compute_perfs,

627 comm_perfs,

628 recompute_perfs,

629 recomm_perfs,

630 debugger=debugger,

631 )

632 logger.info("PerfEst: stage_perfs %s", stage_perfs)

633

634 stage_focused = kwargs.get("stage_focused", None)

635 perf = estimate_perf(

636 cfg, ccfg, stage_perfs, stage_focused=stage_focused, debugger=debugger

637 )

638 perf += estimate_p2p(cfg, ccfg, stage_perfs, debugger=debugger)

639 logger.info("PerfEst: perf %s", perf)

640

641 cache_file = kwargs.get("cache_file")

642 return _finalize_perf(perf, cache_file, debugger, memory) # / cfg.gbs

643

644# TO-DO

645# Fix More Memory

646# Add Context Parallelism

647# Fix PerformanceType.TIME

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

277 statements