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

1# Copyright 2025-2026 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"""Experimental : Comm time""" 

16from copy import copy, deepcopy 

17from typing import Any, NamedTuple 

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

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

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

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

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

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

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

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

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

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

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

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

30 EvalAttn, 

31 EvalFFn, 

32 EvalNorm, 

33) 

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

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

36 get_layer_group_configs, 

37 get_layer_switches_by_position, 

38 get_model_order, 

39 get_table_quantity, 

40) 

41from hyper_parallel.auto_parallel.sapp_nd.nd.common.cp_types import ( 

42 CPCommunicationCost, 

43 CPAlgo, 

44 _resolve_cp_algo, 

45) 

46from hyper_parallel.auto_parallel.sapp_nd.nd.common.arch_hooks import apply_layer_kind, layer_kinds 

47from hyper_parallel.auto_parallel.sapp_nd.nd.common.derive import param_cp 

48from hyper_parallel.auto_parallel.sapp_nd.nd.common.cost_model_preprocess import ( 

49 detect_attention_type, 

50 AttentionType, 

51 compute_kv_dim, 

52 CostModelConfig, 

53) 

54 

55COUNT_OPTIMIZER = False 

56 

57 

58def _cp_resolve_topology(cp, device_per_node, bw_intra, bw_inter): 

59 """Resolve CP topology and effective bandwidth. 

60 

61 Returns: 

62 Tuple of (topology_str, effective_bandwidth). 

63 """ 

64 intra_ranks = min(int(cp), int(device_per_node)) 

65 if cp <= device_per_node: 

66 return "intra-node", bw_intra 

67 if intra_ranks == 1: 

68 return "cross-node", bw_inter 

69 intra_fraction = (intra_ranks - 1) / (cp - 1) 

70 cross_fraction = 1.0 - intra_fraction 

71 bw = intra_fraction * bw_intra + cross_fraction * bw_inter 

72 return "mixed", bw 

73 

74 

75def _cp_comm_zero(ccfg): 

76 """Return a zero CPCommunicationCost for cp <= 1.""" 

77 return CPCommunicationCost( 

78 kv_volume_per_step=0.0, total_kv_volume=0.0, comm_volume=0.0, 

79 ring_steps=0, ring_directions=0, 

80 total_comm_time=0.0, exposed_comm_time=0.0, 

81 overlap_ratio=0.5, effective_bandwidth=0.0, 

82 topology="none", cp_degree=int(ccfg.cp), 

83 seq_len=int(ccfg.s), batch_size=int(ccfg.b), 

84 attention_type=AttentionType.MHA, kv_dim=0, 

85 cp_algo=CPAlgo.COLOSSALAI_CP, 

86 ) 

87 

88 

89class _CPVolumes(NamedTuple): 

90 """Algorithm-specific CP communication volumes of one layer. 

91 

92 ``comm_volume`` is what :func:`cp_traffic` moves, the only field a score 

93 reads; the others are logged. 

94 """ 

95 

96 kv_volume_per_step: float 

97 total_kv_volume: float 

98 comm_volume: float 

99 ring_steps: int 

100 ring_directions: int 

101 

102 

103def _cp_comm_cost_common(ccfg, volumes, attention_type, kv_dim, cp_algo, 

104 topology, effective_bandwidth): 

105 """Build CPCommunicationCost with standard time calculation.""" 

106 overlap_ratio = 0.5 

107 total_comm_time = (volumes.total_kv_volume / (effective_bandwidth * 1e9)) * 1e3 

108 exposed_comm_time = total_comm_time * (1 - overlap_ratio) 

109 return CPCommunicationCost( 

110 kv_volume_per_step=volumes.kv_volume_per_step, 

111 total_kv_volume=volumes.total_kv_volume, 

112 comm_volume=volumes.comm_volume, 

113 ring_steps=volumes.ring_steps, ring_directions=volumes.ring_directions, 

114 total_comm_time=total_comm_time, 

115 exposed_comm_time=exposed_comm_time, 

116 overlap_ratio=overlap_ratio, 

117 effective_bandwidth=effective_bandwidth, 

118 topology=topology, cp_degree=int(ccfg.cp), 

119 seq_len=int(ccfg.s), batch_size=int(ccfg.b), 

120 attention_type=attention_type, kv_dim=int(kv_dim), 

121 cp_algo=cp_algo, 

122 ) 

123 

124 

125def _ulysses_cp_volumes(ccfg): 

126 """Ulysses CP volumes: two All2All over (cp-1)/cp of the local shard.""" 

127 s, b = ccfg.s, ccfg.b 

128 cp = ccfg.cp 

129 t = max(1, ccfg.t) 

130 local_qkv = s * b * (ccfg.a / t) * ccfg.dh * 2 

131 a2a_vol = local_qkv * (cp - 1) / cp 

132 return _CPVolumes(a2a_vol, a2a_vol * 2, cp_traffic(ccfg, CPAlgo.ULYSSES_CP), 0, 2) 

133 

134 

135def _ring_cp_volumes(ccfg, kv_dim, cp_algo): 

136 """Ring CP volumes: cp-1 P2P steps of s/cp tokens of KV, both directions.""" 

137 s, b = ccfg.s, ccfg.b 

138 cp = ccfg.cp 

139 kv_bytes = 4 

140 kv_vol_step = (s / cp) * b * kv_dim * kv_bytes 

141 total_kv = kv_vol_step * (cp - 1) * 2 

142 return _CPVolumes(kv_vol_step, total_kv, cp_traffic(ccfg, cp_algo), int(cp - 1), 2) 

143 

144 

145def cp_traffic(ccfg: CostModelConfig, cp_algo: CPAlgo) -> float: 

146 """The bytes a rank moves a micro-batch for one layer's context parallelism. 

147 

148 Attention exchanges its keys and values, the MLP nothing. colossalai 

149 and hybrid CP all-gather K and V over the sequence in the forward and 

150 reduce-scatter their gradients in the backward, as HyperParallel runs 

151 them (a ring passes the same chunks), each moving (cp - 1) / cp of the 

152 sequence a rank. Ulysses all-to-alls the local query, key, value and 

153 output between sequence and heads, forward and backward. A 

154 linear-attention layer passes its recurrent state to the next rank, 

155 and its gradient back, unless it all-to-alls as Ulysses does. 

156 """ 

157 cp, t = ccfg.cp, max(1, ccfg.t) 

158 ring = (cp - 1) / cp 

159 head = ccfg.dh or ccfg.h / max(1, ccfg.a) 

160 kv_width = compute_kv_dim(ccfg) 

161 if cp_algo == CPAlgo.ULYSSES_CP: 

162 rope = ccfg.dhr if detect_attention_type(ccfg) == AttentionType.MLA else 0 

163 widths = ccfg.a * (head + rope) / t + 2 * kv_width + ccfg.a * head / t 

164 return 2 * ring * ccfg.s / cp * ccfg.b * widths * ccfg.bytes_compute 

165 if ccfg.n_linrec: 

166 state_bytes = 4 

167 return 2 * ccfg.a / t * head * head * state_bytes 

168 return 2 * ring * ccfg.s * ccfg.b * 2 * kv_width * ccfg.bytes_compute 

169 

170 

171def cp_comm_layer_detailed(ccfg: CostModelConfig, ctx: Context = None) -> CPCommunicationCost: 

172 """Estimate CP communication cost with detailed breakdown. 

173 

174 Ring CP (colossalai_cp / hybrid_cp): 

175 Ring P2P in both FW and BW directions. 

176 Each step transfers (s/cp) tokens of KV data. 

177 Total KV volume = kv_volume_per_step * (cp-1) * 2 directions. 

178 

179 Ulysses CP: 

180 All2All in both FW and BW (2 All2All total). 

181 Each All2All: every rank sends (cp-1)/cp of its local shard 

182 and receives the rest from other ranks. 

183 Per-All2All volume = s * b * (a/t) * bytes * (cp-1)/cp (head dims). 

184 Total volume = 2 * per-All2All volume. 

185 

186 *ctx* is not read: the exchange a recompute sends again is priced by 

187 the walk that recomputes it (``_recomputed_comm``). 

188 """ 

189 del ctx 

190 if ccfg.cp <= 1: 

191 return _cp_comm_zero(ccfg) 

192 

193 if ccfg.a <= 0: 

194 raise ValueError(f"Number of attention heads must be positive, got {ccfg.a}") 

195 

196 kv_dim = compute_kv_dim(ccfg) 

197 attention_type = detect_attention_type(ccfg) 

198 cp_algo = _resolve_cp_algo(ccfg) 

199 topology, effective_bandwidth = _cp_resolve_topology( 

200 ccfg.cp, ccfg.device_per_node, ccfg.bw_intra, ccfg.bw_inter) 

201 

202 if cp_algo == CPAlgo.ULYSSES_CP: 

203 volumes = _ulysses_cp_volumes(ccfg) 

204 else: 

205 volumes = _ring_cp_volumes(ccfg, kv_dim, cp_algo) 

206 return _cp_comm_cost_common( 

207 ccfg, volumes, attention_type, kv_dim, cp_algo, topology, effective_bandwidth) 

208 

209 

210def fill_dp_table(cfg, tables): 

211 """DP""" 

212 table_dp = {} 

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

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

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

216 

217 if COUNT_OPTIMIZER: 

218 table_dp["n_attParamCast"] = ( 

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

220 ) 

221 table_dp["n_ffParamCast"] = ( 

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

223 ) 

224 for op in table_dp: 

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

226 

227 table_exp_dp = deepcopy(table_dp) 

228 table_exp_dp["n_ffMM"] = ( 

229 2 

230 * (cfg.n_exp + cfg.n_shared_exp) 

231 * cfg.h 

232 * cfg.hff_exp 

233 / cfg.t 

234 * cfg.bytes_p 

235 ) 

236 tables[Dim.DP] = table_dp 

237 tables["exp_dp"] = table_exp_dp 

238 

239 

240def fill_tp_table(cfg, tables): 

241 """TP""" 

242 table_tp = {} 

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

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

245 

246 for op in table_tp: 

247 table_tp[op] *= cfg.bytes_compute 

248 

249 table_exp_tp = deepcopy(table_tp) 

250 table_exp_tp["n_gather"] = ( 

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

252 ) 

253 tables["tp"] = table_tp 

254 tables["exp_tp"] = table_exp_tp 

255 

256 

257def fill_ep_table(cfg, tables, device_type): 

258 """EP""" 

259 intra_devices = device_type.intra_node_num() 

260 table_ep = {} 

261 inter_node_bias_ep = 1 

262 table_ep["n_ffMM"] = ( 

263 4 

264 * cfg.n_chosen_exp 

265 * cfg.b 

266 * cfg.s 

267 * cfg.h 

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

269 * cfg.cap_fact 

270 * ( 

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

272 + ( 

273 inter_node_bias_ep 

274 * cfg.os_max_shard 

275 / (cfg.ep / intra_devices) 

276 if cfg.ep > intra_devices 

277 else 0 

278 ) 

279 ) 

280 ) 

281 

282 for op in table_ep: 

283 table_ep[op] *= cfg.bytes_compute 

284 tables[Dim.EP] = table_ep 

285 

286 

287def dp_ratio(cfg, device_type): 

288 """formula""" 

289 return ( 

290 0 

291 if cfg.comm_d_non_exp == 0 

292 else 1 

293 - True # overlap_dp, Completely overlap standard DP comm 

294 + ( 

295 1 / 16 

296 if cfg.n_exp == 1 

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

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

299 * (cfg.comm_d_non_exp - 1) 

300 * cfg.os_max_shard 

301 / cfg.d 

302 ) 

303 

304 

305def comm_embed_ouput(cfg): 

306 """ "formula""" 

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

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

309 return comm_embed, comm_output 

310 

311 

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

313 """FW + BW""" 

314 param = { 

315 "cfg": args[0], 

316 "ccfg": args[1], 

317 "stages": args[2], 

318 "device_type": args[3], 

319 "with_recomp": kwargs.get( 

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

321 ), 

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

323 } 

324 

325 param["tables"] = {} 

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

327 

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

329 

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

331 

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

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

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

335 ) 

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

337 

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

339 

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

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

342 

343 lccfgs = get_layer_group_configs(param["cfg"]) 

344 logger.info(lccfgs) 

345 param["layer_count"] = 0 

346 param["idx_lccfg"] = 0 

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

348 # ignores comm recomp, to improve 

349 for stage in param["stages"]: 

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

351 for chunk in stage: 

352 for layer in chunk: 

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

354 estimate_op_bulk_comm_layer( 

355 param, 

356 lccfgs, 

357 layer=layer, 

358 layer_count=param["layer_count"], 

359 idx_lccfg=param["idx_lccfg"], 

360 ) 

361 ) 

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

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

364 comm[dim] = estimate_comm_score( 

365 param["cfg"], 

366 comm[dim], 

367 dim, 

368 overlap=ov, 

369 device=param["device_type"], 

370 ) 

371 

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

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

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

375 

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

377 logger.info("A3 ratio") 

378 comm[Dim.TP] /= 3 

379 

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

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

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

383 

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

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

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

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

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

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

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

391 

392 res = [] 

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

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

395 

396 return res 

397 

398 

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

400 """for estimate_op_bulk_comm""" 

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

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

403 return kwargs["layer_count"] 

404 

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

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

407 if cfg.dc_kv != 0: # Deepseek 

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

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

410 lccfg, 

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

412 LayerType.NOT_REC_LAYER, 

413 kwargs["param"]["with_recomp"], 

414 ) 

415 return kwargs["layer_count"] 

416 

417 if ( 

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

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

420 ): 

421 kwargs["layer_count"] = 0 

422 kwargs["idx_lccfg"] += 1 

423 

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

425 is_moe_layer = lccfg.n_exp > 1 

426 

427 if is_moe_layer: 

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

429 lccfg, 

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

431 kwargs["layer"], 

432 kwargs["param"]["with_recomp"], 

433 ) 

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

435 lccfg, 

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

437 kwargs["layer"], 

438 kwargs["param"]["with_recomp"], 

439 ) 

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

441 lccfg, 

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

443 kwargs["layer"], 

444 kwargs["param"]["with_recomp"], 

445 ) 

446 else: 

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

448 lccfg, 

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

450 kwargs["layer"], 

451 kwargs["param"]["with_recomp"], 

452 ) 

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

454 lccfg, 

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

456 kwargs["layer"], 

457 kwargs["param"]["with_recomp"], 

458 ) 

459 

460 kwargs["layer_count"] += 1 

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

462 

463 

464def prepare_context(): 

465 """context object""" 

466 ctx = Context() 

467 ctx.attn_num_p = EvalAttn.num_params_attn 

468 ctx.ffn_num_p = EvalFFn.num_params_ffn 

469 # A MoE layer's experts, counted as the memory path's eval config counts them. 

470 ctx.ffn_routed_num_p = EvalFFn.num_params_routed_expert 

471 ctx.ffn_shared_num_p = EvalFFn.num_params_shared_expert 

472 ctx.ffn_router_num_p = EvalFFn.num_params_router 

473 ctx.norm_num_p = EvalNorm.num_params_norm 

474 

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

476 EvalHead.num_params_embed, None, None 

477 ) 

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

479 EvalTail.num_params_output, None, None 

480 ) 

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

482 EvalBody.num_params_layer, None, None 

483 ) 

484 ctx.enable_accu_log = False 

485 return ctx 

486 

487 

488def _recomputed_comm(cfg, ctx, layer): 

489 """TP, EP and CP volume a recomputed layer transfers again. 

490 

491 It is the communication whose buffers the layer's memory no longer keeps: 

492 all of it for a fully recomputed layer, and for a selective one what its 

493 switches drop, which the memory model's own terms give as the plain 

494 volume less the selective one. Parameter traffic is not recomputed. Of 

495 CP's traffic, a recompute runs the forward's exchange again, the half 

496 the gather switch keeps in a selective layer. 

497 """ 

498 def _volumes(node): 

499 ctx.current_node = node 

500 kept = EvalUtils.rec_coeff(node == LayerType.SEL_REC_LAYER, EvalUtils.switch(cfg, ctx, "gather")) 

501 return ( 

502 EvalLayerComm.tp_comm_layer(cfg, ctx, 1), 

503 EvalLayerComm.ep_comm_layer(cfg, ctx, 1), 

504 cp_comm_layer_detailed(cfg, ctx).comm_volume / 2 * kept, 

505 ) 

506 

507 kept = ctx.current_node 

508 try: 

509 plain = _volumes(LayerType.NOT_REC_LAYER) 

510 if layer == LayerType.FULL_REC_LAYER: 

511 return plain 

512 selective = _volumes(LayerType.SEL_REC_LAYER) 

513 return tuple(whole - left for whole, left in zip(plain, selective)) 

514 finally: 

515 ctx.current_node = kept 

516 

517 

518def fsdp_root_parts(cfg: Any, layer: Any) -> tuple: 

519 """The root's table, the embedding's or the output layer's, as FSDP holds it. 

520 

521 The parts are as :meth:`EvalLayerComm.fsdp_layer_parts` gives a layer's; 

522 a table the output layer shares is its. 

523 """ 

524 if layer == LayerType.EMBEDDING_LAYER: 

525 if EvalHead.shares_output_table(cfg): 

526 return () 

527 tp = cfg.shard_embed / max(1, cfg.gather_embed or 1) 

528 return ((EvalHead.num_params_embed(cfg, None), tp, cfg.shard_embed, cfg.d * tp),) 

529 return ((EvalTail.num_params_output(cfg, None), cfg.t, cfg.shard_p_os_non_exp_partial, cfg.d * cfg.cp * cfg.t),) 

530 

531 

532def _fsdp_rounds(cfg) -> tuple: 

533 """How many times a micro-batch FSDP gathers a layer's and the root's parameters, and all-reduces. 

534 

535 Resharded after its forward and its backward, a layer is gathered for 

536 both, and the root, kept through its backward, once; kept gathered, each 

537 is gathered once a step. The copies of a shard all-reduce once a step, 

538 after the reduce-scatters, where an FSDP group shards the parameters, 

539 and every micro-batch where none does. 

540 """ 

541 micro = max(1, cfg.m) 

542 reshards = bool(getattr(cfg, "reshards", False)) 

543 # Sharded beyond TP and whatever CP adds to the width (param_cp): with 

544 # TP times CP, a HyperParallel shard no wider than CP read as none. 

545 shards = cfg.shard_p_os_non_exp_partial > cfg.t * param_cp(cfg) 

546 return (2 if reshards else 1 / micro), (1 if reshards else 1 / micro), (1 / micro if shards else 1) 

547 

548 

549def _accumulate_stage_comm(param, stage, stage_id): 

550 """Sum the per-layer DP, TP, EP and CP communication volumes of one stage. 

551 

552 With ``param["with_recomp"]`` or ``param["both"]``, the walk also keeps 

553 the same sums with the volume each recomputed layer's recompute transfers 

554 again added, the way the compute estimate counts a recomputed op twice. 

555 Both sums take every term in the same order as a walk keeping one of 

556 them, so each comes out as that walk's would. 

557 

558 A body layer is priced on the walk's config, in its own kind; the 

559 embedding and the output layer on the model's config, as its family left 

560 it, whichever layer the walk reached last. Under FSDP (grad_shard_as_params), 

561 a layer's DP volume, and the root's, is the traffic its collectives 

562 move a micro-batch (:meth:`EvalLayerComm.fsdp_traffic`). 

563 

564 Returns: 

565 ``(comm, again, reduced)``: the volumes by dimension; the same with 

566 each recomputed layer's traffic added again, or ``None`` where the 

567 walk keeps no such sums; and the part of the DP volume the copies of 

568 a shard all-reduce (:meth:`EvalLayerComm.fsdp_reduce_traffic`). 

569 """ 

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

571 again = dict(comm) if param["with_recomp"] or param["both"] else None 

572 sums = (comm,) if again is None else (comm, again) 

573 reduced = 0.0 

574 fsdp = bool(getattr(param["cfg"], "grad_shard_as_params", False)) 

575 layer_gathers, root_gathers, reduces = _fsdp_rounds(param["cfg"]) if fsdp else (0, 0, 0) 

576 for chunk_id, chunk in enumerate(stage): 

577 for lay_id, layer in enumerate(chunk): 

578 param["ctx"].current_node = layer 

579 position = (stage_id, chunk_id, lay_id) 

580 is_body = layer not in [LayerType.EMBEDDING_LAYER, LayerType.OUTPUT_LAYER] 

581 cfg = param["walk"] if is_body else param["cfg"] 

582 # A layer's own recompute switches, where the config states several. 

583 param["ctx"].switches = param["switches"].get(position) if is_body else None 

584 if is_body and position in param["kinds"]: 

585 kind = param["kinds"][position] 

586 if kind is not None: 

587 apply_layer_kind(cfg, kind) 

588 logger.info("is layer moe ? %s", cfg.n_exp > 1) 

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

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

591 if fsdp: 

592 parts = EvalLayerComm.fsdp_layer_parts(cfg, param["ctx"]) 

593 _add(sums, Dim.DP, EvalLayerComm.fsdp_traffic(cfg, parts, layer_gathers, reduces)) 

594 reduced += EvalLayerComm.fsdp_reduce_traffic(cfg, parts, reduces) 

595 else: 

596 _add(sums, Dim.DP, EvalLayerComm.dp_comm_layer(cfg, param["ctx"])) 

597 elif fsdp and not is_body: 

598 parts = fsdp_root_parts(cfg, layer) 

599 _add(sums, Dim.DP, EvalLayerComm.fsdp_traffic(cfg, parts, root_gathers, reduces)) 

600 reduced += EvalLayerComm.fsdp_reduce_traffic(cfg, parts, reduces) 

601 

602 _add(sums, Dim.TP, EvalLayerComm.tp_comm_layer( 

603 cfg, param["ctx"], 1 

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

605 _add(sums, Dim.EP, EvalLayerComm.ep_comm_layer( 

606 cfg, param["ctx"], 1 

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

608 if is_body: 

609 _add(sums, Dim.CP, cp_comm_layer_detailed( 

610 cfg, param["ctx"] 

611 ).comm_volume) 

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

613 if again is not None and layer in (LayerType.FULL_REC_LAYER, LayerType.SEL_REC_LAYER): 

614 tp_again, ep_again, cp_again = _recomputed_comm(cfg, param["ctx"], layer) 

615 again[Dim.TP] += tp_again 

616 again[Dim.EP] += ep_again 

617 again[Dim.CP] += cp_again 

618 return comm, again, reduced 

619 

620 

621def _add(sums, dim, volume): 

622 """Add *volume* to dimension *dim* of each of *sums*.""" 

623 for comm in sums: 

624 comm[dim] += volume 

625 

626 

627def _stage_comm_time(param, comm): 

628 """Turn one stage's volumes into the walk's units, in place: the time model, the nodes, the overlap, A3.""" 

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

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

631 comm[dim] = estimate_comm_score( 

632 param["cfg"], 

633 comm[dim], 

634 dim, 

635 overlap=ov, 

636 device=param["device_type"], 

637 ) 

638 

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

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

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

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

643 

644 # Transitional overlap correction. 

645 # The search runs the FLOP path, which has no other overlap 

646 # modeling; these factors are the only overlap correction on that 

647 # path. The TIME path's estimate_comm_score(overlap=...) call 

648 # above is zeroed, so this is the single source of overlap for 

649 # both paths. 

650 # Defaults (dp=0.9, tp=0.5) are MindFormers-validated overlap, not 

651 # test hacks: they made the model match real MindFormers step times. 

652 # Re-validating for the hyper-parallel target is a follow-up. 

653 # Follow-up: source from hardware, fix estimate_comm_score's dim 

654 # list and add latency, then fold this into estimate_comm_score. 

655 comm[Dim.DP] *= (1 - param["cfg"].comm_dp_overlap) 

656 comm[Dim.TP] *= (1 - param["cfg"].comm_tp_overlap) 

657 

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

659 logger.info("A3 ratio") 

660 comm[Dim.DP] /= 2 

661 comm[Dim.TP] /= 2 

662 comm[Dim.EP] /= 2 

663 comm[Dim.CP] /= 2 

664 

665 

666def _stage_totals(comms): 

667 """Each stage's volume over the four dimensions, from per-dimension lists.""" 

668 return [c + comms[Dim.DP][i] + comms[Dim.EP][i] + comms[Dim.CP][i] for i, c in enumerate(comms[Dim.TP])] 

669 

670 

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

672 """Each stage's communication, from the memory model's volumes. 

673 

674 With ``both``, one walk returns ``(plain, recomputed)``: each stage's 

675 communication, then the same with each recomputed layer's traffic added 

676 again, which ``with_recomp`` alone returns. The debugger records the 

677 plain parts. 

678 """ 

679 

680 param = { 

681 "cfg": args[0], 

682 "ccfg": args[1], 

683 "stages": args[2], 

684 "device_type": args[3], 

685 } 

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

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

688 ) 

689 param["with_recomp"] = kwargs.get( 

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

691 ) 

692 param["both"] = kwargs.get("both", False) 

693 param["ctx"] = prepare_context() 

694 # The layers' kinds go on a copy: the model's config leaves the walk as 

695 # it entered it, for the next walk and every read after. 

696 param["walk"] = copy(param["cfg"]) 

697 

698 # Each layer's kind, in model order; layers past the stack get no entry, 

699 # so no kind and no DP term. 

700 param["kinds"] = dict(zip(get_model_order(param["cfg"], param["stages"]), layer_kinds(param["cfg"]))) 

701 param["switches"] = get_layer_switches_by_position(param["cfg"], param["stages"]) 

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

703 recomms = {Dim.DP: [], Dim.TP: [], Dim.EP: [], Dim.CP: []} 

704 # Each stage's all-reduce, the share of its DP time the copies of a shard 

705 # spend reducing it: estimate_pipeline takes it out of DP_COMM. 

706 reduces = [] 

707 for stage_id, stage in enumerate(param["stages"]): 

708 plain, again, reduced = _accumulate_stage_comm(param, stage, stage_id) 

709 comm = again if param["with_recomp"] and not param["both"] else plain 

710 reduce_share = reduced / comm[Dim.DP] if comm[Dim.DP] else 0.0 

711 

712 _stage_comm_time(param, comm) 

713 for dim, values in comms.items():