Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / perf_estimation / comm_time.py: 91%
274 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-04 05:18 +0800
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 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)
37from hyper_parallel.auto_parallel.sapp_nd.nd.common.cp_types import (
38 CPCommunicationCost,
39 CPAlgo,
40 _resolve_cp_algo,
41)
42from hyper_parallel.auto_parallel.sapp_nd.nd.common.cost_model_preprocess import (
43 detect_attention_type,
44 AttentionType,
45 compute_kv_dim,
46 CostModelConfig,
47)
49COUNT_OPTIMIZER = False
52def _cp_resolve_topology(cp, device_per_node, bw_intra, bw_inter):
53 """Resolve CP topology and effective bandwidth.
55 Returns:
56 Tuple of (topology_str, effective_bandwidth).
57 """
58 intra_ranks = min(int(cp), int(device_per_node))
59 if cp <= device_per_node:
60 return "intra-node", bw_intra
61 if intra_ranks == 1:
62 return "cross-node", bw_inter
63 intra_fraction = (intra_ranks - 1) / (cp - 1)
64 cross_fraction = 1.0 - intra_fraction
65 bw = intra_fraction * bw_intra + cross_fraction * bw_inter
66 return "mixed", bw
69def _cp_comm_zero(ccfg):
70 """Return a zero CPCommunicationCost for cp <= 1."""
71 return CPCommunicationCost(
72 kv_volume_per_step=0.0, total_kv_volume=0.0, comm_volume=0.0,
73 ring_steps=0, ring_directions=0,
74 total_comm_time=0.0, exposed_comm_time=0.0,
75 overlap_ratio=0.5, effective_bandwidth=0.0,
76 topology="none", cp_degree=int(ccfg.cp),
77 seq_len=int(ccfg.s), batch_size=int(ccfg.b),
78 attention_type=AttentionType.MHA, kv_dim=0,
79 cp_algo=CPAlgo.COLOSSALAI_CP,
80 )
83def _cp_comm_cost_common(volume_per_step, total_kv_volume, comm_volume,
84 ring_steps, ring_directions, cp, s, b,
85 attention_type, kv_dim, cp_algo, topology,
86 effective_bandwidth):
87 """Build CPCommunicationCost with standard time calculation."""
88 overlap_ratio = 0.5
89 total_comm_time = (total_kv_volume / (effective_bandwidth * 1e9)) * 1e3
90 exposed_comm_time = total_comm_time * (1 - overlap_ratio)
91 return CPCommunicationCost(
92 kv_volume_per_step=volume_per_step,
93 total_kv_volume=total_kv_volume,
94 comm_volume=comm_volume,
95 ring_steps=ring_steps, ring_directions=ring_directions,
96 total_comm_time=total_comm_time,
97 exposed_comm_time=exposed_comm_time,
98 overlap_ratio=overlap_ratio,
99 effective_bandwidth=effective_bandwidth,
100 topology=topology, cp_degree=int(cp),
101 seq_len=int(s), batch_size=int(b),
102 attention_type=attention_type, kv_dim=int(kv_dim),
103 cp_algo=cp_algo,
104 )
107def cp_comm_layer_detailed(ccfg: CostModelConfig, ctx: Context = None) -> CPCommunicationCost:
108 """Estimate CP communication cost with detailed breakdown.
110 Ring CP (colossalai_cp / hybrid_cp):
111 Ring P2P in both FW and BW directions.
112 Each step transfers (s/cp) tokens of KV data.
113 Total KV volume = kv_volume_per_step * (cp-1) * 2 directions.
115 Ulysses CP:
116 All2All in both FW and BW (2 All2All total).
117 Each All2All: every rank sends (cp-1)/cp of its local shard
118 and receives the rest from other ranks.
119 Per-All2All volume = s * b * (a/t) * bytes * (cp-1)/cp (head dims).
120 Total volume = 2 * per-All2All volume.
121 """
122 if ccfg.cp <= 1:
123 return _cp_comm_zero(ccfg)
125 s, b = ccfg.s, ccfg.b
126 cp = ccfg.cp
127 t = max(1, ccfg.t)
129 if ccfg.a <= 0:
130 raise ValueError(f"Number of attention heads must be positive, got {ccfg.a}")
132 kv_dim = compute_kv_dim(ccfg)
133 attention_type = detect_attention_type(ccfg)
134 cp_algo = _resolve_cp_algo(ccfg)
135 topology, effective_bandwidth = _cp_resolve_topology(
136 cp, ccfg.device_per_node, ccfg.bw_intra, ccfg.bw_inter)
138 # rec_factor: recompute coefficient matching old cp_comm_non_exp
139 rec_layer = (ctx.current_node == LayerType.SEL_REC_LAYER) if ctx else False
140 rec_op_gather = getattr(getattr(ccfg, 'rec_op', None), 'gather', 0)
141 rec_factor = (int(not rec_layer) | rec_op_gather) * int(ccfg.p == 1)
143 if cp_algo == CPAlgo.ULYSSES_CP:
144 local_qkv = s * b * (ccfg.a / t) * ccfg.dh * 2
145 a2a_vol = local_qkv * (cp - 1) / cp
146 # comm_volume: same weighted-unit as dp/tp/ep
147 # Ulysses attention coeff = 0.5*rec_factor + 0.5
148 ulysses_attn_coeff = 0.5 * rec_factor + 0.5
149 comm_vol = (
150 ccfg.comm_cp * 2 * s * b
151 * (ulysses_attn_coeff * ccfg.n_attMM * ccfg.h
152 + ccfg.n_ffMM * ccfg.hff)
153 / t
154 )
155 return _cp_comm_cost_common(
156 a2a_vol, a2a_vol * 2, comm_vol, 0, 2, cp, s, b,
157 attention_type, kv_dim, cp_algo, topology, effective_bandwidth)
159 kv_bytes = 4
160 kv_vol_step = (s / cp) * b * kv_dim * kv_bytes
161 total_kv = kv_vol_step * (cp - 1) * 2
162 # comm_volume: same weighted-unit as dp/tp/ep
163 # Ring attention coeff = 2*0.5*rec_factor + 0.5 (extra /cp from (s/cp)^2)
164 ring_attn_coeff = 2 * 0.5 * rec_factor + 0.5
165 comm_vol = (
166 ccfg.comm_cp * 2 * s * b
167 * (ring_attn_coeff * ccfg.n_attMM * ccfg.h
168 + ccfg.n_ffMM * ccfg.hff)
169 / t
170 )
171 return _cp_comm_cost_common(
172 kv_vol_step, total_kv, comm_vol, int(cp - 1), 2, cp, s, b,
173 attention_type, kv_dim, cp_algo, topology, effective_bandwidth)
176def fill_dp_table(cfg, tables):
177 """DP"""
178 table_dp = {}
179 table_dp["n_attMM"] = cfg.h * cfg.h / cfg.t
180 table_dp["n_ffMM"] = cfg.h * cfg.hff / cfg.t
181 table_dp["n_normOp"] = 2 * cfg.h / cfg.sp
183 if COUNT_OPTIMIZER:
184 table_dp["n_attParamCast"] = (
185 11 * cfg.h * cfg.h / (cfg.d if cfg.has_op else 1)
186 )
187 table_dp["n_ffParamCast"] = (
188 11 * cfg.h * cfg.hff / (cfg.d if cfg.has_op else 1)
189 )
190 for op in table_dp:
191 table_dp[op] *= cfg.bytes_norm if op == "n_normOp" else cfg.bytes_p
193 table_exp_dp = deepcopy(table_dp)
194 table_exp_dp["n_ffMM"] = (
195 2
196 * (cfg.n_exp + cfg.n_shared_exp)
197 * cfg.h
198 * cfg.hff_exp
199 / cfg.t
200 * cfg.bytes_p
201 )
202 tables[Dim.DP] = table_dp
203 tables["exp_dp"] = table_exp_dp
206def fill_tp_table(cfg, tables):
207 """TP"""
208 table_tp = {}
209 high_tp_bias = 11 / 16 if cfg.t >= 8 else 1 # Fix this
210 table_tp["n_gather"] = cfg.b * cfg.s * cfg.h * high_tp_bias
212 for op in table_tp:
213 table_tp[op] *= cfg.bytes_compute
215 table_exp_tp = deepcopy(table_tp)
216 table_exp_tp["n_gather"] = (
217 cfg.b * cfg.s * cfg.h * 1.5 * (cfg.ep / cfg.d) * cfg.bytes_compute
218 )
219 tables["tp"] = table_tp
220 tables["exp_tp"] = table_exp_tp
223def fill_ep_table(cfg, tables, device_type):
224 """EP"""
225 intra_devices = device_type.intra_node_num()
226 table_ep = {}
227 inter_node_bias_ep = 1
228 table_ep["n_ffMM"] = (
229 4
230 * cfg.n_chosen_exp
231 * cfg.b
232 * cfg.s
233 * cfg.h
234 * (max(4, cfg.os_max_shard) / cfg.t)
235 * cfg.cap_fact
236 * (
237 cfg.os_max_shard / min(intra_devices, cfg.ep)
238 + (
239 inter_node_bias_ep
240 * cfg.os_max_shard
241 / (cfg.ep / intra_devices)
242 if cfg.ep > intra_devices
243 else 0
244 )
245 )
246 )
248 for op in table_ep:
249 table_ep[op] *= cfg.bytes_compute
250 tables[Dim.EP] = table_ep
253def dp_ratio(cfg, device_type):
254 """formula"""
255 return (
256 0
257 if cfg.comm_d_non_exp == 0
258 else 1
259 - True # overlap_dp, Completely overlap standard DP comm
260 + (
261 1 / 16
262 if cfg.n_exp == 1
263 else 1 / max(1, cfg.ep / device_type.intra_node_num()) / 1.25
264 ) # overlap_op, Bias in overlapping OP comm (todo:make it dynamic too)
265 * (cfg.comm_d_non_exp - 1)
266 * cfg.os_max_shard
267 / cfg.d
268 )
271def comm_embed_ouput(cfg):
272 """ "formula"""
273 comm_embed = cfg.bytes_compute * cfg.h * cfg.v / cfg.shard_embed
274 comm_output = cfg.h * cfg.v / cfg.t
275 return comm_embed, comm_output
278def estimate_op_bulk_comm(*args, **kwargs):
279 """FW + BW"""
280 param = {
281 "cfg": args[0],
282 "ccfg": args[1],
283 "stages": args[2],
284 "device_type": args[3],
285 "with_recomp": kwargs.get(
286 "with_recomp", args[4] if len(args) > 4 else False
287 ),
288 "debugger": kwargs.get("debugger", args[5] if len(args) > 5 else None),
289 }
291 param["tables"] = {}
292 fill_dp_table(param["cfg"], param["tables"])
294 param['dp_ratio'] = dp_ratio(param['cfg'], param['device_type'])
296 param["comm_embed"], param["comm_output"] = comm_embed_ouput(param["cfg"])
298 if param["cfg"].dc_kv != 0: # Deepseek
299 param["comm_output"] += param["cfg"].h * (
300 2 * param["cfg"].h + param["cfg"].v
301 )
302 param["comm_output"] *= param["cfg"].n_mtp
304 param["comm_output"] *= param["cfg"].bytes_p
306 fill_tp_table(param["cfg"], param["tables"])
307 fill_ep_table(param["cfg"], param["tables"], param["device_type"])
309 lccfgs = get_layer_custom_configs(param["cfg"])
310 logger.info(lccfgs)
311 param["layer_count"] = 0
312 param["idx_lccfg"] = 0
313 comms = {Dim.DP: [], Dim.TP: [], Dim.EP: []}
314 # ignores comm recomp, to improve
315 for stage in param["stages"]:
316 comm = {Dim.DP: 0.0, Dim.TP: 0.0, Dim.EP: 0.0}
317 for chunk in stage:
318 for layer in chunk:
319 param["layer_count"], param["idx_lccfg"] = (
320 estimate_op_bulk_comm_layer(
321 param,
322 lccfgs,
323 layer=layer,
324 layer_count=param["layer_count"],
325 idx_lccfg=param["idx_lccfg"],
326 )
327 )
328 if param["ccfg"].ttype == PerformanceType.TIME:
329 for dim, ov in zip([Dim.DP, Dim.TP, Dim.DP], [0.0, 0.0, 0.0]):
330 comm[dim] = estimate_comm_score(
331 param["cfg"],
332 comm[dim],
333 dim,
334 overlap=ov,
335 device=param["device_type"],
336 )
338 comm[Dim.DP] *= param["dp_ratio"]
339 comm[Dim.TP] *= param["cfg"].comm_t
340 comm[Dim.EP] *= param["cfg"].comm_ep
342 if param["device_type"].name == "A3":
343 logger.info("A3 ratio")
344 comm[Dim.TP] /= 3
346 comms[Dim.DP].append(comm[Dim.DP])
347 comms[Dim.TP].append(comm[Dim.TP])
348 comms[Dim.EP].append(comm[Dim.EP])
350 if param["debugger"] and param["debugger"].is_enabled():
351 logger.info("DP_COMM = %s", comms[Dim.DP])
352 logger.info("MP_COMM = %s", comms[Dim.TP])
353 logger.info("EP_COMM = %s", comms[Dim.EP])
354 param["debugger"].info[PerfParts.DP_COMM] = comms[Dim.DP]
355 param["debugger"].info[PerfParts.MP_COMM] = comms[Dim.TP]
356 param["debugger"].info[PerfParts.EP_COMM] = comms[Dim.EP]
358 res = []
359 for i, c in enumerate(comms[Dim.TP]):
360 res.append(comms[Dim.DP][i] + c + comms[Dim.EP][i])
362 return res
365def estimate_op_bulk_comm_layer(cfg, lccfgs, **kwargs):
366 """for estimate_op_bulk_comm"""
367 if kwargs["layer"] == LayerType.EMBEDDING_LAYER:
368 kwargs["comm"][Dim.DP] += kwargs["param"]["comm_embed"]
369 return kwargs["layer_count"]
371 if kwargs["layer"] == LayerType.OUTPUT_LAYER:
372 kwargs["comm"][Dim.DP] += kwargs["param"]["comm_output"]
373 if cfg.dc_kv != 0: # Deepseek
374 lccfg = lccfgs[kwargs["idx_lccfg"]][0]
375 kwargs["comm"][Dim.TP] += cfg.n_mtp * get_table_quantity(
376 lccfg,
377 kwargs["param"]["tables"]["exp_tp"],
378 LayerType.NOT_REC_LAYER,
379 kwargs["param"]["with_recomp"],
380 )
381 return kwargs["layer_count"]
383 if (
384 kwargs["idx_lccfg"] + 1 < len(lccfgs)
385 and lccfgs[kwargs["idx_lccfg"]][1] == kwargs["layer_count"]
386 ):
387 kwargs["layer_count"] = 0
388 kwargs["idx_lccfg"] += 1
390 lccfg = lccfgs[kwargs["idx_lccfg"]][0]
391 is_moe_layer = lccfg.n_exp > 1
393 if is_moe_layer:
394 kwargs["comm"][Dim.DP] += get_table_quantity(
395 lccfg,
396 kwargs["param"]["tables"]["exp_dp"],
397 kwargs["layer"],
398 kwargs["param"]["with_recomp"],
399 )
400 kwargs["comm"][Dim.TP] += get_table_quantity(
401 lccfg,
402 kwargs["param"]["tables"]["exp_tp"],
403 kwargs["layer"],
404 kwargs["param"]["with_recomp"],
405 )
406 kwargs["comm"][Dim.EP] += get_table_quantity(
407 lccfg,
408 kwargs["param"]["tables"][Dim.EP],
409 kwargs["layer"],
410 kwargs["param"]["with_recomp"],
411 )
412 else:
413 kwargs["comm"][Dim.DP] += get_table_quantity(
414 lccfg,
415 kwargs["param"]["tables"][Dim.DP],
416 kwargs["layer"],
417 kwargs["param"]["with_recomp"],
418 )
419 kwargs["comm"][Dim.TP] += get_table_quantity(
420 lccfg,
421 kwargs["param"]["tables"]["tp"],
422 kwargs["layer"],
423 kwargs["param"]["with_recomp"],
424 )
426 kwargs["layer_count"] += 1
427 return kwargs["layer_count"], kwargs["idx_lccfg"]
430def prepare_context():
431 """context object"""
432 ctx = Context()
433 ctx.attn_num_p = EvalAttn.num_params_attn
434 ctx.ffn_num_p = EvalFFn.num_params_ffn
435 ctx.norm_num_p = EvalNorm.num_params_norm
437 ctx.node_eval[LayerType.EMBEDDING_LAYER] = NodeEval(
438 EvalHead.num_params_embed, None, None
439 )
440 ctx.node_eval[LayerType.OUTPUT_LAYER] = NodeEval(
441 EvalTail.num_params_output, None, None
442 )
443 ctx.node_eval[LayerType.NOT_REC_LAYER] = NodeEval(
444 EvalBody.num_params_layer, None, None
445 )
446 ctx.enable_accu_log = False
447 return ctx
450def estimate_from_mem_comm(*args, **kwargs):
451 """For memory estimation"""
453 param = {
454 "cfg": args[0],
455 "ccfg": args[1],
456 "stages": args[2],
457 "device_type": args[3],
458 }
459 param["debugger"] = kwargs.get(
460 "debugger", args[5] if len(args) > 5 else None
461 )
462 param["ctx"] = prepare_context()
464 # For layer type
465 param["flatten"] = sum(
466 [[f[1]] * f[0] for f in param["cfg"].layer_custom_config], []
467 )
468 comms = {Dim.DP: [], Dim.TP: [], Dim.EP: [], Dim.CP: []}
469 for stage in param["stages"]:
470 comm = {Dim.DP: 0.0, Dim.TP: 0.0, Dim.EP: 0.0, Dim.CP: 0.0}
471 for chunk in stage:
472 for layer in chunk:
473 param["ctx"].current_node = layer
474 if (
475 layer
476 not in [LayerType.EMBEDDING_LAYER, LayerType.OUTPUT_LAYER]
477 and param["flatten"]
478 ):
479 custom_fun = param["flatten"].pop(0)
480 if custom_fun:
481 custom_fun(param["cfg"])
482 logger.info("is layer moe ? %s", param["cfg"].n_exp > 1)
483 param["ctx"].current_node = LayerType.NOT_REC_LAYER
484 logger.info("param ctx %s", param["ctx"])
485 comm[Dim.DP] += EvalLayerComm.dp_comm_layer(param["cfg"], param["ctx"])
487 comm[Dim.TP] += EvalLayerComm.tp_comm_layer(
488 param["cfg"], param["ctx"], 1
489 ) # / 4 #* (param["cfg"].t - 1)
490 comm[Dim.EP] += EvalLayerComm.ep_comm_layer(
491 param["cfg"], param["ctx"], 1
492 ) # * param["cfg"].ep
493 comm[Dim.CP] += cp_comm_layer_detailed(
494 param["cfg"], param["ctx"]
495 ).comm_volume
496 # min(device_type.level_bound_number[0], param["cfg"].ep)
497 # comm_cp += EvalLayerComm.cp_comm_layer
498 # (param["cfg"], param["ctx"])
502 if param["ccfg"].ttype == PerformanceType.TIME:
503 for dim, ov in zip([Dim.DP, Dim.TP, Dim.CP], [0.0, 0.0, 0.0]):
504 comm[dim] = estimate_comm_score(
505 param["cfg"],
506 comm[dim],
507 dim,
508 overlap=ov,
509 device=param["device_type"],
510 )
512 dev_per_node = param["device_type"].level_bound_number[0]
513 comm[Dim.TP] *= max(1, param["cfg"].t // dev_per_node)
514 comm[Dim.EP] *= max(1, param["cfg"].ep // dev_per_node)
515 comm[Dim.CP] *= max(1, param["cfg"].cp // dev_per_node)
517 # Transitional overlap correction.
518 # The search runs the FLOP path, which has no other overlap
519 # modeling; these factors are the only overlap correction on that
520 # path. The TIME path's estimate_comm_score(overlap=...) call
521 # above is zeroed, so this is the single source of overlap for
522 # both paths.
523 # Defaults (dp=0.9, tp=0.5) are MindFormers-validated overlap, not
524 # test hacks: they made the model match real MindFormers step times.
525 # Re-validating for the hyper-parallel target is a follow-up.
526 # Follow-up: source from hardware, fix estimate_comm_score's dim
527 # list and add latency, then fold this into estimate_comm_score.
528 comm[Dim.DP] *= (1 - param["cfg"].comm_dp_overlap)
529 comm[Dim.TP] *= (1 - param["cfg"].comm_tp_overlap)
531 if param["device_type"].name == "A3":
532 logger.info("A3 ratio")
533 comm[Dim.DP] /= 2
534 comm[Dim.TP] /= 2
535 comm[Dim.EP] /= 2
536 comm[Dim.CP] /= 2
538 comms[Dim.DP].append(comm[Dim.DP])
539 comms[Dim.TP].append(comm[Dim.TP])
540 comms[Dim.EP].append(comm[Dim.EP])
541 comms[Dim.CP].append(comm[Dim.CP])
543 if param["debugger"] and param["debugger"].is_enabled():
544 logger.info("DP_COMM = %s", comms[Dim.DP])
545 logger.info("MP_COMM = %s", comms[Dim.TP])
546 logger.info("EP_COMM = %s", comms[Dim.EP])
547 logger.info("CP_COMM = %s", comms[Dim.CP])
548 param["debugger"].info[PerfParts.DP_COMM] = comms[Dim.DP]
549 param["debugger"].info[PerfParts.MP_COMM] = comms[Dim.TP]
550 param["debugger"].info[PerfParts.EP_COMM] = comms[Dim.EP]
551 param["debugger"].info[PerfParts.CP_COMM] = comms[Dim.CP]
552 if param["cfg"].cp > 1:
553 cp_comm_details = cp_comm_layer_detailed(param["cfg"], param["ctx"])
554 param["debugger"].info["CP_KV_VOLUME"] = cp_comm_details.total_kv_volume
555 param["debugger"].info["CP_EXPOSED_TIME"] = cp_comm_details.exposed_comm_time
556 param["debugger"].info["CP_TOPOLOGY"] = cp_comm_details.topology
557 param["debugger"].info["CP_BANDWIDTH"] = cp_comm_details.effective_bandwidth
559 res = []
560 for i, c in enumerate(comms[Dim.TP]):
561 res += [c + comms[Dim.DP][i] + comms[Dim.EP][i] + comms[Dim.CP][i]]
563 return res
566def estimate_comm(*args, **kwargs):
567 """wrapper"""
568 cfg, ccfg, stages, device_type = args[0], args[1], args[2], args[3]
569 with_recomp = kwargs.get(
570 "with_recomp", args[4] if len(args) > 4 else False
571 )
572 debugger = kwargs.get("debugger", args[5] if len(args) > 5 else None)
573 # return estimate_op_bulk_comm(cfg, ccfg, stages,
574 # device_type=device_type, with_recomp=with_recomp,
575 # debugger=debugger)
576 return estimate_from_mem_comm(
577 cfg,
578 ccfg,
579 stages,
580 device_type,
581 with_recomp=with_recomp,
582 debugger=debugger,
583 )
586def level_efficiency(level):
587 """to improve for Ascend A2"""
588 if level == NetworkLevel.NODE:
589 return 0.7
590 if level == NetworkLevel.CLUSTER:
591 return 0.9
592 raise ValueError
595def level_bandwidth(level):
596 """to improve for Ascend A2"""
597 if level == NetworkLevel.NODE:
598 return 300
599 if level == NetworkLevel.CLUSTER:
600 return 25
601 raise ValueError
604def level_latency(level):
605 """to improve for Ascend A2"""
606 if level == NetworkLevel.NODE:
607 return 0.00001
608 if level == NetworkLevel.CLUSTER:
609 return 0.00002
610 raise ValueError
613def comm_throughput(level):
614 """formula"""
615 eff = level_efficiency(level)
616 bw = level_bandwidth(level)
617 return bw * eff
620def estimate_comm_size_time(_, comm_size, level):
621 """formula"""
622 th = comm_throughput(level)
623 lat = level_latency(level)
624 return lat + comm_size / th
627def estimate_comm_score(
628 cfg, comm_volume, dim, overlap=0.0, device=Hard.device_map["A2"]
629):
630 """score assignment"""
631 assignment = device.level_assign(dp=cfg.d, tp=cfg.t, cp=cfg.cp, pp=cfg.p)
632 score = 0
633 for level in range(device.levels):
634 # intra_comm = comm_volume * (1-overlap)
635 # * (assignment[dim][0]-1) / device.intra_node_bw
636 score += (
637 comm_volume
638 * (1 - overlap)
639 * (
640 (assignment[dim][level] - 1)
641 * device.devices_below_level(level)
642 / device.level_bandwidth[level]
643 )
644 )
645 return score