Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / perf_estimation / estimate.py: 97%
428 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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"""performance estimation"""
16import functools
17import json
18import math
19import os
20from copy import deepcopy
21from dataclasses import replace
22from types import MappingProxyType
23from typing import Any, Callable, Dict, NamedTuple, Optional, Tuple
24import numpy as np
26from hyper_parallel.auto_parallel.sapp_nd.nd.logger import logger as nd_logger
27from hyper_parallel.auto_parallel.sapp_nd.nd.logger import perf_logger as logger
28from hyper_parallel.auto_parallel.sapp_nd.nd.common.config import Config
29from hyper_parallel.auto_parallel.sapp_nd.nd.common.layer_type import LayerType
30from hyper_parallel.auto_parallel.sapp_nd.nd.common.cost_model_config import CostModelConfig
31from hyper_parallel.auto_parallel._layer_stack import LayerStack, StackGroup
32from hyper_parallel.auto_parallel._op_profiles import LayerKind
33from hyper_parallel.auto_parallel.sapp_nd.nd.common.arch_hooks import CWrap, check_and_apply_custom_hook
34import hyper_parallel.auto_parallel.sapp_nd.nd.common.hardware as Hard
35from hyper_parallel.auto_parallel.sapp_nd.nd.debug import Debug, PerfParts, RealParts, estimation_in_real_parts
37from hyper_parallel.auto_parallel.sapp_nd.perf_estimation.utils_classes import (
38 RatioType,
39 PerformanceType,
40 P2PCommType,
41 RecType,
42 CustomConfig,
43)
44from hyper_parallel.auto_parallel.sapp_nd.perf_estimation.comm_time import estimate_comm
45from hyper_parallel.auto_parallel.sapp_nd.perf_estimation.getters import (
46 get_layer_configs_by_position,
47 get_layer_switches_by_position,
48 get_table_quantity,
49 mla_weights,
50 selective_shares,
51)
54GENERALIZE_PIPELINE_CALCULATION = False
55MANUAL_P2P_RATIO = 0.002
56BACKWARD_RATIO = 2
57# The chunk the gated delta rule's kernel and Transformers' own
58# implementation run a sequence in.
59GDN_CHUNK = 64
60# What a MoE layer's dispatch costs a rank, per token, per hidden unit, per
61# expert a token routes to and per rank holding experts. The work is the
62# sorting, permuting, zeroing and accumulating around the all-to-all, which
63# no matmul count sees, and it is the one cost that measurably grows with the
64# expert-parallel degree. Over 15 profiled points of Qwen3.5-35B-A3B at 8192
65# tokens on 64 dies, device compute averages 5.85, 6.19, 6.59, 7.58 and 9.43 s
66# at EP 1, 2, 4, 8 and 16, a span of 61%, where this model without the term
67# spans 8%. The increments over EP 1 are 222, 217, 230 and 230 ms a unit of
68# EP, so they follow EP itself rather than EP - 1: a layer sorts and permutes
69# its tokens by expert even where it holds every one of them, and the work
70# grows with the ranks it spreads them over. With the term the estimate runs
71# 3.8 to 9.5% above the measured compute at every one of those degrees, where
72# it was 31% below at EP 16, and the rank correlation over the 15 points rises
73# from +0.51 to +0.64. It does not change which strategy the search picks:
74# what holds that at EP 1 is the gather priced at 370 ms where 204 was
75# measured, an overlap threshold rather than a compute term.
76#
77# EMPIRICAL, and fitted at one shape. It follows the tokens a rank routes, so
78# on the 128-token sweeps it adds 3.2 ms a unit of EP, where the measured
79# compute is flat from EP 4 to EP 64, 196 to 210 ms, and this would add 205 ms
80# by EP 64: at a shape whose step is launch bound, 53 to 67% of it idle,
81# neither this term nor the compute model under it holds. The kernel tables
82# per EP at the training shape name what grows: the routing's sort and index
83# kernels, launched as often at every degree above 1 and each running longer,
84# a cost linear in the tokens the all-to-all hands a rank. This coefficient
85# stands in for that until it is priced from the tokens, so read it as a
86# calibration carrying the measured shape, not as a model of the dispatch. A
87# cluster that has measured its own says so with ``context.moe_dispatch``, as
88# it states the ratios of the parts in a file, and a model whose compute does
89# not grow with EP states 0, which prices no dispatch: DeepSeek V4.1's crop
90# measures 3.3% less compute at EP 16 than at EP 1, where ND with this term
91# prices 60% more.
92MOE_DISPATCH = 1040
93# How far a stage's recorded parts may stray from its time: they are its time
94# taken apart, so only by rounding.
95PARTS_REL_TOL = 1e-9
96# The ratio of a ratios file that scales the compute parts and turns ND's
97# units into milliseconds.
98COMPUTE_RATIO = "COMPUTE"
99# Where a ratio a file leaves out is taken from, as ``nd.ratios.fit_ratios``
100# fills a part it cannot fit: the all-reduce from the DP traffic it was split
101# out of, the bubble from P2P time, which measured data cannot tell it from,
102# and any other part from compute's, which leaves it as ND weighs it.
103RATIO_FALLBACKS = {
104 PerfParts.DP_REDUCE.name: PerfParts.DP_COMM.name,
105 PerfParts.BUBBLE.name: PerfParts.PP_COMM.name,
106}
107# The ratios a file was found without, each named once a process.
108_MISSING_RATIOS = set()
109# The parts a communication walk records on the debugger, one value a stage.
110RECORDED_COMM_PARTS = (
111 PerfParts.DP_COMM,
112 PerfParts.DP_REDUCE,
113 PerfParts.MP_COMM,
114 PerfParts.EP_COMM,
115 PerfParts.CP_COMM,
116)
119def op_table(cfg, attn=None):
120 """op compute load formulas
122 Args:
123 cfg: the cost-model config.
124 attn: a layer group whose attention flavour differs from *cfg*'s.
125 Only the attention entries read it; everything else stays on
126 *cfg*, so the expert scaling applied by the caller is unaffected.
127 """
128 table = {}
129 # cfg.s /= cfg.cp
130 # Q,O cost h x (a*dh) each, K,V cost h x (n_kv*dh); the h*h form is the
131 # head_dim = h/a special case. A fused output gate doubles Q.
132 att = attn if attn is not None else cfg
133 d_h = att.dh or (cfg.h / att.a if att.a else 0)
134 d_q = att.a * d_h
135 table["n_attMM"] = (
136 1.5
137 * cfg.b
138 * cfg.s
139 * cfg.h
140 * (
141 ((2 if getattr(att, "attn_output_gate", False) else 1) + 1) * d_q
142 + 2 * att.n_kv * d_h
143 )
144 )
145 # The chunked gated delta rule, per value head: within a chunk of
146 # GDN_CHUNK tokens, its keys against its keys and its queries, its
147 # solved weights against its values and its decayed keys, and its
148 # scores against its new values; the state read twice and written once
149 # a chunk. Linear in the sequence, where an attention score is
150 # quadratic. The entry exists only for a group that declares the
151 # flavour, so no other config needs to carry the count.
152 n_v, d_k, d_v = (getattr(att, name, 0) for name in ("lin_n_v", "lin_d_k", "lin_d_v"))
153 if n_v * d_k * d_v:
154 table["n_linrec"] = (3 * cfg.b * cfg.s * n_v
155 * (2 * GDN_CHUNK * (3 * d_k + 2 * d_v) + 6 * d_k * d_v))
156 table["n_ffMM"] = 6 * cfg.b * cfg.s * cfg.h * cfg.hff
157 # Every head's queries against every key, as wide as a head's queries
158 # and keys, an MLA head's with its rotary part; then the weights
159 # against the values, as wide as a head's values.
160 d_qk = (getattr(att, "qk_nope_head_dim", None) or d_h) + (getattr(att, "dhr", 0) or 0)
161 table["n_attBMM"] = 3 * cfg.b * cfg.s * cfg.s * att.a * (d_qk + d_h)
162 table["n_softmax"] = 13 * cfg.a * cfg.b * cfg.s * cfg.s
163 table["n_headCast"] = 3 * cfg.a * cfg.b * cfg.s * cfg.s
164 table["n_gather"] = cfg.b * cfg.s * cfg.h * (cfg.t - 1)
165 table["n_ffAct"] = 21 * cfg.b * cfg.s * cfg.hff
167 table["n_normOp"] = 30 * cfg.b * cfg.s * cfg.h * cfg.t / cfg.sp
168 # A QK-norm runs over every head's queries and keys, which TP splits.
169 # The entry exists only for a model that has one.
170 if getattr(cfg, "n_qknorm", 0):
171 table["n_qknorm"] = 30 * cfg.b * cfg.s * (att.a + att.n_kv) * d_h
172 table["n_dropout"] = (
173 3 * cfg.b * cfg.s * max(cfg.a * cfg.s, 3 * cfg.h * cfg.t / cfg.sp)
174 )
175 if cfg.dc_kv != 0: # MLA
176 weights, _ = mla_weights(cfg)
177 table["n_attMM"] = 6 * cfg.b * cfg.s * weights / (getattr(cfg, "n_attMM", 0) or 4)
178 for op in table:
179 table[op] *= cfg.bytes_p / cfg.t / cfg.cp
180 # cfg.s *= cfg.cp
181 return table
184def _flavour_tables(cfg, attn=None):
185 """One (dense, expert) table pair for a given attention flavour.
187 A MoE layer's feed-forward entry prices the experts each token runs,
188 the shared experts, as wide as a routed one, and, spread over the
189 layer's feed-forward matmuls, the router and the shared experts' gate;
190 its activation function's entry the same experts, which run it at
191 their width, where a dense layer's runs it at the dense width. A MoE
192 layer also dispatches its tokens, which the table prices apart
193 (:data:`MOE_DISPATCH`, or the cost the run states, 0 pricing none); a
194 dense layer does not, so only the expert table carries that entry.
195 """
196 base = op_table(cfg, attn)
197 exp = deepcopy(base) # Verify this with MF MoEV2
198 n_ff = getattr(attn if attn is not None else cfg, "n_ffMM", 0) or 3
199 gate = 1 if cfg.n_shared_exp and getattr(cfg, "shared_expert_gate", None) else 0
200 experts = cfg.hff_exp * (max(1, cfg.n_chosen_exp) * cfg.cap_fact + cfg.n_shared_exp)
201 width = experts + (cfg.n_exp + gate) / n_ff
202 exp["n_ffMM"] *= width / cfg.hff
203 exp["n_ffAct"] *= experts / cfg.hff
204 dispatch = getattr(cfg, "moe_dispatch", None)
205 if dispatch is None:
206 dispatch = MOE_DISPATCH
207 exp["n_dispatch"] = (
208 dispatch * cfg.b * cfg.s * cfg.h * max(1, cfg.n_chosen_exp)
209 * max(1, getattr(cfg, "ep", 1) or 1) * cfg.bytes_p / cfg.t / cfg.cp
210 )
211 return base, exp
214def _regular_layer_flop(cfg, ccfg, lcfg, tables, layer, with_recomp, switches=None):
215 """Price one regular layer with its group's config, and its own recompute *switches* where it has them.
217 *tables* holds one table pair per attention flavour and gains one the
218 first time a flavour is priced.
219 """
220 kind = getattr(lcfg, "attn_kind", "full")
221 if kind not in tables:
222 tables[kind] = _flavour_tables(cfg, lcfg)
223 flop = get_table_quantity(
224 lcfg,
225 tables[kind][1] if (lcfg.n_exp > 1) else tables[kind][0],
226 layer,
227 with_recomp,
228 shares=selective_shares(lcfg, layer, switches),
229 switches=switches,
230 )
231 if ccfg.ttype == PerformanceType.TIME:
232 flop = estimate_comp_flop_time(lcfg, flop)
233 return flop
236# Evaluation functions
237def estimate_op_bulk_comp(cfg, ccfg, stages, with_recomp=False, debugger=None, both=False):
238 """FW + BW; with ``both``, ``(plain, recomputed)`` from one walk over the layers"""
239 _ = debugger
240 # Full attention prices from the model's own attention dimensions.
241 tables = {"full": _flavour_tables(cfg)}
242 lccfg_at = get_layer_configs_by_position(cfg, stages)
243 # Each layer's own recompute switches, where the config states several.
244 switches_at = get_layer_switches_by_position(cfg, stages)
246 # One sum a mode, each taking its terms in the order a walk of that mode alone did.
247 modes = (False, True) if both else (with_recomp,)
248 sums = tuple([] for _ in modes)
249 for stage_id, stage in enumerate(stages):
250 for flops in sums:
251 flops += [0]
252 for chunk_id, chunk in enumerate(stage):
253 for lay_id, layer in enumerate(chunk):
254 if layer == LayerType.EMBEDDING_LAYER:
255 continue
257 if layer == LayerType.OUTPUT_LAYER:
258 output = (1 if cfg.dc_kv == 0 else cfg.n_mtp) * (
259 1
260 / 16 # bias_imbalance
261 * 6
262 * cfg.b
263 * cfg.v
264 * cfg.h
265 * cfg.s
266 * cfg.bytes_p
267 / cfg.t
268 )
269 for flops in sums:
270 flops[-1] += output
271 continue
273 position = (stage_id, chunk_id, lay_id)
274 for flops, mode in zip(sums, modes):
275 flops[-1] += _regular_layer_flop(
276 cfg, ccfg, lccfg_at[position], tables, layer, mode, switches_at.get(position)
277 )
279 return sums if both else sums[0]
282def estimate_comp(cfg, ccfg, stages, with_recomp=False, debugger=None, both=False):
283 """wrapper"""
284 return estimate_op_bulk_comp(
285 cfg, ccfg, stages, with_recomp, debugger=debugger, both=both
286 )
289# Experimental : Flop time
292def efficiency(x):
293 """obtained via extrapolation"""
294 eff = min(
295 1.0, max(0.1, 0.00004694 * x**3 + 0.0014 * x**2 - 0.0336 * x + 0.1)
296 )
297 return eff
300def throughput(precision_bytes, flop):
301 """assumes matrix
303 Only PerformanceType.TIME reaches this, and nothing in production asks
304 for it: the live estimate is a relative score whose scale the fitted
305 ratios own. Its precision_bytes**2 * 1e12 is no device's rate,
306 4 TFLOP/s at bf16, and no other rate prices a candidate's parts: the
307 one ND states, a host link's sustained_tflops, only turns an offload
308 copy's seconds into the score's units where no fitted ratio does.
309 """
310 eff = efficiency(flop / (10.0**12))
311 return precision_bytes**2 * (10.0**12) * eff
314def estimate_comp_flop_time(cfg, flop, is_softmax=False):
315 """flop from throughput"""
316 th = throughput(
317 cfg.bytes_softmax if is_softmax else cfg.bytes_compute, flop
318 )
319 return flop / th
322##############################################################
325def get_dynamic_ratio(cfg):
326 """comm/comp"""
327 if cfg.n_exp == 1:
328 return 3 / 2 * (cfg.hff + cfg.s) * (8192 / (cfg.h + cfg.s))
329 return 3 / 2 * (cfg.hff_exp + cfg.s) * (8192 / (cfg.h + cfg.s))
332def _weigh_comm_parts(debugger, comm_w):
333 """Weigh each stage's recorded communication parts as the stage time weighs them.
335 DP_REDUCE, the all-reduce's share of DP_COMM, is weighed where the
336 communication walk recorded it (``estimate_from_mem_comm``).
337 """
338 for p in [PerfParts.DP_COMM, PerfParts.MP_COMM, PerfParts.EP_COMM, PerfParts.CP_COMM]:
339 debugger.info[p] = [
340 comm_w * c for c in debugger.info[p]
341 ]
342 if isinstance(debugger.info[PerfParts.DP_REDUCE], list):
343 debugger.info[PerfParts.DP_REDUCE] = [comm_w * c for c in debugger.info[PerfParts.DP_REDUCE]]
346def estimate_stage(*args, **kwargs):
347 """stage level estimation"""
348 cfg = args[0]
349 ccfg = args[1]
350 compute_perfs = args[2]
351 comm_perfs = args[3]
352 recompute_perfs = args[4]
353 recomm_perfs = args[5]
354 debugger = kwargs.get("debugger", args[6] if len(args) > 6 else None)
355 comp_w = 1
356 comm_w = 1
357 if ccfg.rtype == RatioType.COMM_ONLY:
358 comp_w = 0
359 elif ccfg.rtype == RatioType.COMPUTE_ONLY:
360 comm_w = 0
361 elif ccfg.rtype == RatioType.STATIC:
362 comm_w = 10**4
363 ccfg.static_ratio = comm_w
364 elif ccfg.rtype == RatioType.DYNAMIC:
365 comm_w = get_dynamic_ratio(cfg)
366 ccfg.dynamic_ratio = comm_w
367 perf = [
368 comp_w * compute_perfs[i] + comm_w * comm_perfs[i]
369 for i in range(len(compute_perfs))
370 ]
371 logger.info("ratio = %s", comm_w)
372 # ignores comm recomp, to improve
373 re_perf = [
374 (
375 max(0, comp_w * (recompute_perfs[i] - compute_perfs[i]))
376 + max(0, comm_w * (recomm_perfs[i] - comm_perfs[i]))
377 )
378 / (1 + BACKWARD_RATIO)
379 for i in range(len(compute_perfs))
380 ]
382 if debugger and debugger.is_enabled():
383 _weigh_comm_parts(debugger, comm_w)
384 debugger.info[PerfParts.FW_COMPUTE] = [
385 comp_w * comp / (1 + BACKWARD_RATIO) for comp in compute_perfs
386 ]
387 debugger.info[PerfParts.BW_COMPUTE] = [
388 fw * BACKWARD_RATIO for fw in debugger.info[PerfParts.FW_COMPUTE]
389 ]
390 debugger.info[PerfParts.RECOMPUTE] = re_perf
392 return [perf[i] + re_perf[i] for i in range(len(perf))]
393 #penalty_fn(stage)
394 #return stage
397def _check_stage_parts(time_sum, stage_time, stage):
398 """Say so when the parts recorded for a stage do not add up to its time.
400 The parts are the stage's time taken apart, so they agree with it to
401 rounding. A gap means a walk recorded less, or more, than the time
402 counts, and the bubble, what is left of the pipeline's time once the
403 parts are taken out, silently takes the difference.
404 """
405 if math.isclose(time_sum, stage_time, rel_tol=PARTS_REL_TOL, abs_tol=PARTS_REL_TOL):
406 return
407 nd_logger.error(
408 "the parts recorded for stage %d sum to %.6E where its time is %.6E (%+.3f%%), "
409 "and the bubble takes the difference",
410 stage,
411 time_sum,
412 stage_time,
413 100 * (time_sum - stage_time) / stage_time if stage_time else 0.0,
414 )
417def estimate_pipeline(cfg, stage_perfs, stage_focused=None, debugger=None):
418 """pipeline level estimation"""
419 logger.info("stage_perfs = %s", stage_perfs)
420 straggler_time = max(stage_perfs)
421 sum_time = sum(stage_perfs)
422 last_straggler_idx = cfg.p - 1 - np.argmax(stage_perfs[::-1])
423 logger.info(
424 "straggler estim is %s and its stage is %s",
425 straggler_time,
426 last_straggler_idx,
427 )
429 non_steady_perf = 0
430 steady_perf = 0
431 if cfg.p == 1:
432 if len(stage_perfs) != 1:
433 raise ValueError("Expected exactly one stage performance")
434 steady_perf = sum_time * cfg.m
435 elif cfg.vp == 1:
436 non_steady_perf = sum_time
437 if GENERALIZE_PIPELINE_CALCULATION:
438 last_idx = last_straggler_idx + 1
439 steady_perf = (
440 cfg.m - cfg.p + last_straggler_idx
441 ) * straggler_time + sum(stage_perfs[last_idx:])
442 else:
443 steady_perf = (cfg.m - 1) * straggler_time
444 else:
445 less_extra = cfg.p * (cfg.vp - 1)
446 # big_extra = less_extra + cfg.p
448 # we assume that times of all micro-blocks in one vp chunk are the same
449 straggler_time /= cfg.vp
450 sum_time /= cfg.vp
452 # more_memory has more micro blocks in warm-up but it does
453 # not matter since they will overlap with steady phase
454 non_steady_perf = sum_time + less_extra * straggler_time
456 # more_memory is a boost to performance and a nerf to memory
457 # if cfg.vp_less_memory or True:
458 # steady_perf = (cfg.m * cfg.vp - less_extra - 1) * straggler_time
459 # else:
460 # steady_perf = (cfg.m * cfg.vp - big_extra - 1) * straggler_time
461 steady_perf = (cfg.m * cfg.vp - less_extra - 1) * straggler_time
462 straggler_time *= cfg.vp
463 sum_time *= cfg.vp
465 pipeline_perf = non_steady_perf + steady_perf
466 logger.info(
467 "pipeline_perf = non_steady_perf(%.2E) + steady_perf(%.2E)",
468 non_steady_perf,
469 steady_perf,
470 )
472 if stage_focused is not None:
473 last_straggler_idx = stage_focused
474 if debugger and debugger.is_enabled():
475 time_sum = 0
476 for k in [
477 PerfParts.DP_COMM,
478 PerfParts.MP_COMM,
479 PerfParts.EP_COMM,
480 PerfParts.CP_COMM,
481 PerfParts.FW_COMPUTE,
482 PerfParts.BW_COMPUTE,
483 PerfParts.RECOMPUTE,
484 ]:
485 debugger.info[k] = (
486 debugger.info[k][last_straggler_idx] * cfg.m
487 ) # / cfg.vp
488 logger.info(
489 "time_sum += debugger[%s] = %.2E",
490 k,
491 debugger.info[k],
492 )
493 time_sum += debugger.info[k]
495 # The bubble stays the pipeline's time less the parts, so the parts
496 # and the bubble still sum to the total the ratios are fitted on.
497 _check_stage_parts(time_sum, stage_perfs[last_straggler_idx] * cfg.m, last_straggler_idx)
498 logger.info(
499 "straggler time = %.2E. %s x stragglers = %.2E",
500 straggler_time,
501 cfg.m,
502 straggler_time * cfg.m,
503 )
505 bubble = pipeline_perf - time_sum
506 logger.info(
507 "bubble(%.2E) = pipeline_perf(%.2E) - time_sum(%.2E)",
508 bubble,
509 pipeline_perf,
510 time_sum,
511 )
512 debugger.info[PerfParts.BUBBLE] = bubble
513 # DP_COMM timed the stage's whole DP traffic, the time and the bubble
514 # above with it; its all-reduce is DP_REDUCE's, the rest FSDP's.
515 reduced = debugger.info[PerfParts.DP_REDUCE]
516 if isinstance(reduced, list):
517 debugger.info[PerfParts.DP_REDUCE] = reduced[last_straggler_idx] * cfg.m
518 debugger.info[PerfParts.DP_COMM] -= debugger.info[PerfParts.DP_REDUCE]
519 return pipeline_perf
522def estimate_p2p_comm(cfg, straggler, ratio=MANUAL_P2P_RATIO, debugger=None):
523 """pipeline comm"""
524 nb_send_recv = 0
525 if cfg.vp == 1:
526 nb_send_recv = (
527 0
528 if cfg.p == 1
529 else (
530 4 * cfg.m
531 if cfg.p == 2
532 else 4 * cfg.p * cfg.m + 4 * cfg.p * cfg.p - 14 * cfg.p
533 )
534 )
535 else:
536 nb_send_recv = (
537 0
538 if cfg.p == 1
539 else (
540 8 * cfg.m * cfg.vp - 4 * cfg.m
541 if cfg.p == 2
542 else (
543 16 * cfg.m * cfg.vp + 12
544 if cfg.p == 4
545 else 4 * cfg.p * cfg.m * cfg.vp
546 + 4 * cfg.p * cfg.p
547 - 13 * cfg.p
548 )
549 )
550 )
551 pp_comm = ratio * nb_send_recv / cfg.p * straggler / cfg.sp
552 if debugger and debugger.is_enabled():
553 debugger.info[PerfParts.PP_COMM] = pp_comm
555 return pp_comm
558def estimate_perf(cfg, _, stage_perfs, stage_focused=None, debugger=None):
559 """wrapper"""
560 return estimate_pipeline(cfg, stage_perfs, stage_focused=stage_focused, debugger=debugger)
563def estimate_p2p(cfg, ccfg, stage_perfs, debugger=None):
564 """wrapper"""
565 if ccfg.ptype != P2PCommType.MANUAL:
566 p2p = 0
567 else:
568 p2p = estimate_p2p_comm(cfg, max(stage_perfs), debugger=debugger)
569 if debugger and debugger.is_enabled():
570 debugger.info[PerfParts.PP_COMM] = p2p
571 return p2p
574def _one_layer_of(stack: Optional[LayerStack], kind: Optional[LayerKind]) -> Optional[LayerStack]:
575 """The model's stack, reduced to one layer of *kind*.
577 Its other groups stay, with no layer, so the estimates apply the layer's
578 kind exactly when they would for the whole model: to a copy of the config
579 for the compute, whose op tables come from the model's own config, and in
580 place for the communication. A layer with no kind has no stack.
581 """
582 if kind is None or stack is None:
583 return None
584 return replace(stack, groups=(StackGroup(kind, 1),) + tuple(replace(group, count=0) for group in stack.groups))
587class PlainTimes(NamedTuple):
588 """What a layer costs before any recompute, the same for every option of its kind.
590 Attributes:
591 compute: The compute estimate without recompute.
592 comm: The communication estimate without recompute.
593 config: The config the communication estimate leaves, which the stage
594 estimate reads.
595 """
597 compute: float
598 comm: float
599 config: CostModelConfig
602def _layer_alone(
603 cfg: CostModelConfig, kind: Optional[LayerKind], layer_type: LayerType, switches: Optional[Dict[str, int]]
604) -> Tuple[CostModelConfig, list]:
605 """One *layer_type* layer of *kind*, alone in the first stage of a fresh copy of *cfg*.
607 A fresh copy each time: the estimates leave state on the config they are
608 given.
609 """
610 single = deepcopy(cfg)
611 if switches is not None:
612 single.rec_op = Config(dict(switches))
613 # The one layer runs *switches*, else the config's rec_op, whatever
614 # the layers of the model run.
615 single.layer_switches = None
616 single.layer_stack = _one_layer_of(single.layer_stack, kind)
617 single.n = single.d * single.t * single.p
618 stages = [[[] for _ in range(single.vp)] for _ in range(single.p)]
619 stages[0][0] = [layer_type]
620 return single, stages
623def plain_layer_times(
624 cfg: CostModelConfig,
625 kind: Optional[LayerKind],
626 layer_type: LayerType,
627 device_type: Any,
628 ccfg: Optional[CustomConfig] = None,
629) -> PlainTimes:
630 """The compute and communication of one layer of *kind* alone, without recompute.
632 They do not depend on which recompute option the layer runs, so every
633 option of a kind can share them; the embedding and the output layer have
634 their own.
635 """
636 ccfg = ccfg if ccfg is not None else CustomConfig()
637 single, stages = _layer_alone(cfg, kind, layer_type, None)
638 compute = estimate_comp(single, ccfg, stages)[0]
639 grouped, stages = _layer_alone(cfg, kind, layer_type, None)
640 comm = estimate_comm(grouped, ccfg, stages, device_type)[0]
641 return PlainTimes(compute, comm, grouped)
644def estimate_layer_times(
645 cfg: CostModelConfig,
646 kind: Optional[LayerKind],
647 layer_type: LayerType,
648 device_type: Any,
649 ccfg: Optional[CustomConfig] = None,
650 switches: Optional[Dict[str, int]] = None,
651 plain: Optional[PlainTimes] = None,
652) -> Tuple[float, float]:
653 """Forward and backward time of one layer, priced as a stage prices it.
655 Runs the compute, communication and stage estimates of
656 :func:`estimate_performance` on one layer of *kind*, alone in the first
657 stage of a copy of *cfg*. For a stage whose layers are all of one kind,
658 the sum of their times is what the search charges the stage before
659 pipeline bubbles and point-to-point communication. The search prices each
660 layer on a config the previous layers' kinds have been applied to, so a
661 stage that mixes kinds, or holds the embedding or output layer, can
662 differ by a few percent; here every layer is priced on its own kind,
663 wherever it sits.
665 Args:
666 cfg: The model's cost-model config, before any layer kind is applied.
667 kind: The layer's kind; ``None`` for the embedding and output, and
668 for a layer priced on the config as it stands.
669 layer_type: How the layer runs: plain, selective or full recompute,
670 embedding or output.
671 device_type: The device the communication is priced on.
672 ccfg: Estimator options; the search's defaults when omitted.
673 switches: The recompute switches a selective layer runs with, 1 to
674 keep an op and 0 to recompute it; the config's own when omitted.
675 plain: The layer's :func:`plain_layer_times`, when already priced.
677 Returns:
678 ``(forward, backward)``, *backward* including the layer's recompute.
679 """
680 ccfg = ccfg if ccfg is not None else CustomConfig()
681 if plain is None:
682 plain = plain_layer_times(cfg, kind, layer_type, device_type, ccfg)
683 comp, comm, grouped = plain
684 # A layer that recomputes nothing runs nothing again.
685 recomputes = layer_type in (LayerType.FULL_REC_LAYER, LayerType.SEL_REC_LAYER)
686 recomp = comp
687 if recomputes and ccfg.retype in {RecType.COMPUTE_ONLY, RecType.WITH}:
688 single, stages = _layer_alone(cfg, kind, layer_type, switches)
689 recomp = estimate_comp(single, ccfg, stages, with_recomp=True)[0]
690 recomm = comm
691 if recomputes and ccfg.retype in {RecType.COMM_ONLY, RecType.WITH}:
692 fresh, fresh_stages = _layer_alone(cfg, kind, layer_type, switches)
693 recomm = estimate_comm(fresh, ccfg, fresh_stages, device_type, with_recomp=True)[0]
695 # The stage estimate weighs the config the communication estimate leaves.
696 once = estimate_stage(grouped, ccfg, [comp], [comm], [comp], [comm])[0]
697 total = estimate_stage(grouped, ccfg, [comp], [comm], [recomp], [recomm])[0]
698 forward = once / (1 + BACKWARD_RATIO)
699 return forward, total - forward
702class LayerTimes:
703 """Prices the layers of a pipeline balancer's layer description.
705 The memory backbone calls it for every layer it describes, with the config
706 it walks and the layer's kind. A config is copied the first time it is
707 seen, at its embedding layer, before the walk applies any layer kind to
708 it, and each kind, layer type and set of switches is priced once. What a
709 kind costs without recompute is priced once for all its options.
710 """
712 def __init__(self, device_type: Any, ccfg: Optional[CustomConfig] = None) -> None:
713 """Price on *device_type*, with the estimator options *ccfg*."""
714 self._device_type = device_type
715 self._ccfg = ccfg
716 self._base = {}
717 self._times = {}
718 self._plain = {}
720 def __call__(
721 self,
722 cfg: CostModelConfig,
723 kind: Optional[LayerKind],
724 layer_type: LayerType,
725 switches: Optional[Dict[str, int]] = None,
726 ) -> Tuple[float, float]:
727 """``(forward, backward)`` of a *layer_type* layer of *kind*, run with *switches*."""
728 if id(cfg) not in self._base:
729 self._base[id(cfg)] = deepcopy(cfg)
730 key = (id(cfg), kind, layer_type, None if switches is None else tuple(sorted(switches.items())))
731 if key not in self._times:
732 # Every body option of a kind shares its plain times; the
733 # embedding and the output layer have their own.
734 role = layer_type if layer_type in (LayerType.EMBEDDING_LAYER, LayerType.OUTPUT_LAYER) else None
735 plain_key = (id(cfg), kind, role)
736 if plain_key not in self._plain:
737 self._plain[plain_key] = plain_layer_times(
738 self._base[id(cfg)], kind, layer_type, self._device_type, self._ccfg
739 )
740 self._times[key] = estimate_layer_times(
741 self._base[id(cfg)], kind, layer_type, self._device_type, self._ccfg, switches,
742 plain=self._plain[plain_key],
743 )
744 return self._times[key]
747def _part_ratio(coeffs, key):
748 """The ratio a ratios file states for *key*, or the one fit_ratios would give it.
750 A file ``fit_ratios`` writes states every key, but one edited by hand,
751 or fitted by a tree older than a part, may lack some. A missing ratio
752 falls back as ``fit_ratios`` fills a part it cannot fit (``RATIO_FALLBACKS``),
753 and each missing key is named once. Compute's ratio is the unit every
754 other one falls back on, so a file without it is refused.
756 Raises:
757 ValueError: When the file states no COMPUTE ratio.
758 """
759 ratio = coeffs.get(key)
760 if ratio is not None:
761 return ratio
762 if key == COMPUTE_RATIO:
763 raise ValueError(
764 "the ratios file states no COMPUTE ratio, the one that turns ND's "
765 "units into milliseconds and that every missing ratio falls back on"
766 )
767 fallback = RATIO_FALLBACKS.get(key, COMPUTE_RATIO)
768 if key not in _MISSING_RATIOS:
769 _MISSING_RATIOS.add(key)
770 nd_logger.error("the ratios file states no %s ratio: that part takes %s's", key, fallback)
771 return _part_ratio(coeffs, fallback)
774def apply_regression_coefficients(coeffs, debugger, old_perf):
775 """
776 applies the coefficients present in regression's cache_file
777 """
778 compute_ratio = _part_ratio(coeffs, COMPUTE_RATIO)
779 for part, raw in list(debugger.info.items()):
780 if part in (PerfParts.TOTAL, PerfParts.MEMORY):
781 continue
782 if part in (PerfParts.FW_COMPUTE,
783 PerfParts.BW_COMPUTE,
784 PerfParts.RECOMPUTE):
785 ratio = compute_ratio
786 else:
787 ratio = _part_ratio(coeffs, part.name)
788 new_val = 0.0 if raw == 0.0 else raw * ratio
789 debugger.info[part] = new_val
791 max_idx = max(p.value for p in PerfParts) -1
792 estimations = [0.0] * max_idx
793 for part in PerfParts:
794 if part in (PerfParts.TOTAL, PerfParts.MEMORY):
795 continue
796 estimations[part.value - 1] = debugger.info.get(part) or 0.0
798 real_buckets = {rp: [] for rp in RealParts}
799 real_buckets = estimation_in_real_parts(real_buckets, estimations, old_perf)
801 perf = (
802 real_buckets[RealParts.COMP][-1]
803 + real_buckets[RealParts.DP_WAIT][-1]
804 + real_buckets[RealParts.MP_WAIT][-1]
805 + real_buckets[RealParts.EP_WAIT][-1]
806 + real_buckets[RealParts.CP_WAIT][-1]
807 + real_buckets[RealParts.PP_WAIT][-1]
808 )
809 debugger.info[PerfParts.TOTAL] = perf
810 return perf
813class _EstimateArgs(NamedTuple):
814 """Resolved inputs of estimate_performance, unpackable as a tuple."""
816 cfg: CostModelConfig
817 stages: Optional[list]
818 extra_custom_func: Optional[Callable]
819 ccfg: Any
820 debugger: Any
821 device_type: Any
822 memory: Any
825def _resolve_estimate_args(args, kwargs):
826 """Resolve positional/keyword inputs for estimate_performance.
828 Returns:
829 _EstimateArgs of (cfg, stages, extra_custom_func, ccfg, debugger,
830 device_type, memory).
831 """
832 cfg_input = args[0]
833 # A copy: the estimate applies layer kinds to its config in place, and the
834 # caller's config is the one the next estimate starts from.
835 cfg = (
836 deepcopy(cfg_input)
837 if isinstance(cfg_input, CostModelConfig)
838 else CostModelConfig(cfg_input)
839 )
840 stages = kwargs.get("stages", args[1] if len(args) > 1 else None)
841 extra_custom_func = kwargs.get(
842 "extra_custom_func", args[2] if len(args) > 2 else None
843 )
844 ccfg = kwargs.get("ccfg", args[3] if len(args) > 3 else CustomConfig())
845 debugger = kwargs.get("debugger", args[4] if len(args) > 4 else None)
846 device_type = kwargs.get(
847 "device_type", args[5] if len(args) > 5 else Hard.device_map["A2"]
848 )
849 memory = kwargs.get("memory", args[6] if len(args) > 6 else None)
850 return _EstimateArgs(cfg, stages, extra_custom_func, ccfg, debugger, device_type, memory)
853def _save_recompute(stage_perfs, savings, debugger):
854 """Stage times less the recompute time each stage saves.
856 Args:
857 stage_perfs: Each stage's time, its layers recomputed as the config
858 says.
859 savings: The time each stage's layers save with the recompute
860 options chosen for them instead; ``None`` when none are.
861 debugger: Records the recompute part of each stage, which is where
862 the time is saved.
864 Returns:
865 Each stage's time.
866 """
867 if savings is None:
868 return stage_perfs
869 if len(savings) != len(stage_perfs):
870 raise ValueError(f"{len(savings)} stage savings for {len(stage_perfs)} stages")
871 if debugger and debugger.is_enabled():
872 debugger.info[PerfParts.RECOMPUTE] = [
873 recompute - saved for recompute, saved in zip(debugger.info[PerfParts.RECOMPUTE], savings)
874 ]
875 return [perf - saved for perf, saved in zip(stage_perfs, savings)]
878def _user_hook(cfg: Any) -> Optional[Callable]:
879 """The hook a unimodal config's hook class gives it, as the memory estimate applies it."""
880 hooks = getattr(cfg, "hooks_dict", None)
881 if not hooks or getattr(cfg, "multimodal", False):
882 return None
883 return next(iter(hooks.values()))
886def _ratios(cache_file):
887 """The ratios *cache_file* states, parsed again only when the file changes.
889 A search applies them to every candidate it orders: read and parsed for
890 each one, they cost 0.7 s on the Qwen3.5 example's 7216 candidates.
891 """
892 stat = os.stat(cache_file)
893 return _parsed_ratios(cache_file, stat.st_mtime_ns, stat.st_size)
896@functools.lru_cache(maxsize=8)
897def _parsed_ratios(cache_file, mtime_ns, size): # pylint: disable=unused-argument
898 """The ratios file as one read-only mapping; *mtime_ns* and *size* only key the cache."""
899 with open(cache_file, 'r', encoding='utf-8') as f:
900 return MappingProxyType(json.load(f))
903def _finalize_perf(perf, cache_file, debugger, memory):
904 """Apply cached regression coefficients and record debug info.
906 Returns:
907 The final performance value.
908 """
909 coeffs = None
910 cache = False
911 if cache_file is not None:
912 coeffs = _ratios(cache_file)
913 cache = True
915 if debugger and debugger.is_enabled():
916 if cache:
917 perf = apply_regression_coefficients(coeffs, debugger, perf)
918 debugger.info[PerfParts.TOTAL] = perf
919 if memory is not None:
920 debugger.info[PerfParts.MEMORY] = memory
921 return perf
924def _stage_parts(cfg, ccfg, stages, device_type, debugger):
925 """One config's per-stage compute, recompute, communication and recomputed communication."""
926 # Only the plain walks record parts. A stage's time takes what the
927 # recompute walks add as RECOMPUTE alone; recorded as communication too,
928 # it made the parts outgrow the time, and the bubble, their difference,
929 # went negative.
930 # Where the recompute is priced, one walk prices both modes, each sum
931 # taking its terms in the order a walk of that mode alone did.
932 if ccfg.retype in {RecType.COMPUTE_ONLY, RecType.WITH}:
933 compute_perfs, recompute_perfs = estimate_comp(cfg, ccfg, stages, debugger=debugger, both=True)
934 else:
935 compute_perfs = estimate_comp(cfg, ccfg, stages, with_recomp=False, debugger=debugger)
936 recompute_perfs = [0] * cfg.p
937 if ccfg.retype in {RecType.COMM_ONLY, RecType.WITH}:
938 comm_perfs, recomm_perfs = estimate_comm(cfg, ccfg, stages, device_type, debugger=debugger, both=True)
939 else:
940 comm_perfs = estimate_comm(cfg, ccfg, stages, device_type, with_recomp=False, debugger=debugger)
941 recomm_perfs = [0] * cfg.p
942 return compute_perfs, recompute_perfs, comm_perfs, recomm_perfs
945def _submodule_parts(cfg, ccfg, device_type, debugger):
946 """A multimodal model's per-stage parts, every submodule's summed.
948 The submodules of a vision-language model run on one pipeline, so a
949 stage's time is the work of the layers on it, whichever submodule they
950 belong to, as the memory model sums what they keep there
951 (``combine_partition_multimodal``). Each is priced on its own
952 partitions and under its own family, on a copy, so no submodule leaves
953 state for the next.
955 Each submodule's walk also records its communication on the debugger,
956 stage by stage, and the record is summed the same way. Kept as the
957 last walk left it, it held the language model's share alone: a vision
958 tower's collectives went into the bubble, the time less the parts, and
959 the parts a ratio is fitted on missed them.
960 """
961 partitions = cfg.generate_partitions_vpp()
962 totals = None
963 recorded = {}
964 for name in cfg.mm_order:
965 sub = deepcopy(cfg.mm_ccfgs[name])
966 check_and_apply_custom_hook(sub)
967 parts = _stage_parts(sub, ccfg, partitions[name], device_type, debugger)
968 totals = parts if totals is None else tuple(
969 [left + right for left, right in zip(*pair)] for pair in zip(totals, parts)
970 )
971 if debugger and debugger.is_enabled():
972 _add_recorded_comm(recorded, debugger.info)
973 if recorded:
974 debugger.info.update(recorded)
975 return totals
978def _add_recorded_comm(recorded, info):
979 """Add the communication a walk recorded, one value a stage, to *recorded*."""
980 for part in RECORDED_COMM_PARTS:
981 stages = info.get(part)
982 if not isinstance(stages, list):
983 continue
984 recorded[part] = (
985 [left + right for left, right in zip(recorded[part], stages)]
986 if part in recorded
987 else list(stages)
988 )
991# performance estimation
992def estimate_performance(*args, **kwargs):
993 """main estimation
995 With ``stage_savings``, each stage's time is lowered by the recompute
996 time its layers save under the options chosen for them, before the
997 pipeline is timed.
998 """
999 (
1000 cfg,
1001 stages,
1002 extra_custom_func,
1003 ccfg,
1004 debugger,
1005 device_type,
1006 memory,
1007 ) = _resolve_estimate_args(args, kwargs)
1008 cache_file = kwargs.get("cache_file")
1009 if cache_file is not None and not (debugger and debugger.is_enabled()):
1010 # The ratios rescale the parts a debugger records, so the parts are
1011 # recorded whether or not anyone asked to see them: a search that
1012 # writes no debug output must not rank on the uncalibrated score.
1013 debugger = Debug(None, info_type=PerfParts)
1015 # Process custom model config: the config's own hook when its hook class
1016 # gives it one, as the memory estimate applies it, else its family.
1017 user_hook = _user_hook(cfg)
1018 if extra_custom_func:
1019 extra_custom_func(cfg)
1020 elif user_hook is not None:
1021 user_hook(CWrap(cfg))
1022 else:
1023 logger.info("auto applying custom model config")
1024 check_and_apply_custom_hook(cfg)
1026 # Process partition generation
1027 if not stages:
1028 logger.info("stage partitions are generated")
1029 stages = cfg.generate_partitions_vpp()
1031 cfg.n = cfg.d * cfg.t * cfg.p
1033 logger.debug(
1034 "perf_model: DP = %d, TP = %d, EP = %d, PP = %d, MB = %d",
1035 cfg.d,
1036 cfg.t,
1037 cfg.ep,
1038 cfg.p,
1039 cfg.m,
1040 )
1042 logger.info("%s", cfg)
1043 logger.info(stages)
1044 logger.info(ccfg)
1046 if getattr(cfg, "multimodal", False):
1047 compute_perfs, recompute_perfs, comm_perfs, recomm_perfs = _submodule_parts(
1048 cfg, ccfg, device_type, debugger
1049 )
1050 timed = cfg.mm_ccfgs[getattr(cfg, "mm_main", None) or cfg.mm_order[-1]]
1051 timed.n = cfg.n
1052 else:
1053 compute_perfs, recompute_perfs, comm_perfs, recomm_perfs = _stage_parts(
1054 cfg, ccfg, stages, device_type, debugger
1055 )
1056 timed = cfg
1057 logger.info("PerfEst: comm_perfs %s", comm_perfs)
1059 stage_perfs = estimate_stage(
1060 timed,
1061 ccfg,
1062 compute_perfs,
1063 comm_perfs,
1064 recompute_perfs,
1065 recomm_perfs,
1066 debugger=debugger,
1067 )
1068 stage_perfs = _save_recompute(stage_perfs, kwargs.get("stage_savings"), debugger)
1069 logger.info("PerfEst: stage_perfs %s", stage_perfs)
1071 stage_focused = kwargs.get("stage_focused", None)
1072 perf = estimate_perf(
1073 timed, ccfg, stage_perfs, stage_focused=stage_focused, debugger=debugger
1074 )
1075 perf += estimate_p2p(timed, ccfg, stage_perfs, debugger=debugger)
1076 logger.info("PerfEst: perf %s", perf)
1078 return _finalize_perf(perf, cache_file, debugger, memory) # / cfg.gbs
1080# TO-DO
1081# Fix More Memory
1082# Add Context Parallelism