« 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"""Communication volume submodule"""
16from __future__ import annotations
17from typing import TYPE_CHECKING
18from hyper_parallel.auto_parallel.sapp_nd.nd.common.layer_type import LayerType
19from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.utils import EvalUtils
20from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.logger import logger
21from hyper_parallel.auto_parallel.sapp_nd.nd.common.cp_types import (
22 CPAlgo,
23 _resolve_cp_algo,
24)
25from hyper_parallel.auto_parallel.sapp_nd.nd.common.cost_model_preprocess import (
26 compute_kv_dim,
27)
29if TYPE_CHECKING:
30 from hyper_parallel.auto_parallel.sapp_nd.nd.common.cost_model_preprocess import CostModelConfig
31 from hyper_parallel.auto_parallel.sapp_nd.memory_estimation._context import Context
34class EvalLayerComm:
35 """Communication volume formulas class"""
37 @staticmethod
38 def dp_comm_non_exp(ccfg: CostModelConfig, ctx: Context) -> float:
39 """DP/OP comm for non-expert parameters"""
40 non_exp, _, _ = ctx.eval.num_p(ccfg, ctx)
41 dp_comm_non_exp = 0
42 # Non expert ZeRO LvL 2
43 if ccfg.comm_d_non_exp == 2:
44 dp_comm_non_exp += non_exp / (ccfg.cp * ccfg.t)
45 dp_comm_non_exp += non_exp / ccfg.t
46 # Non expert ZeRO LvL 3
47 if ccfg.comm_d_non_exp == 3:
48 dp_comm_non_exp += non_exp / ccfg.t
49 return dp_comm_non_exp
51 @staticmethod
52 def dp_comm_exp(ccfg: CostModelConfig, ctx: Context) -> float:
53 """DP/OP comm for expert parameters"""
54 _, routed, shared = ctx.eval.num_p(ccfg, ctx)
55 exp_param_size = routed + shared
56 if exp_param_size == 0:
57 return 0
58 dp_comm_exp = 0
59 # Expert ZeRO LvL 2
60 if ccfg.comm_d_exp == 2:
61 dp_comm_exp += exp_param_size / (ccfg.cp * ccfg.t_exp * ccfg.ep)
62 dp_comm_exp += exp_param_size / max(ccfg.ep, ccfg.t_exp)
63 # Expert ZeRO LvL 3
64 if ccfg.comm_d_exp == 3:
65 dp_comm_exp += exp_param_size / (ccfg.cp * ccfg.t_exp * ccfg.ep)
66 return dp_comm_exp
68 @staticmethod
69 def dp_comm_layer(ccfg: CostModelConfig, ctx: Context) -> float:
70 """DP/OP comm sum"""
71 non_exp = EvalLayerComm.dp_comm_non_exp(ccfg, ctx)
72 exp = EvalLayerComm.dp_comm_exp(ccfg, ctx)
73 return non_exp + exp
75 @staticmethod
76 def fsdp_layer_parts(ccfg: CostModelConfig, ctx: Context) -> tuple:
77 """A layer's parameters as FSDP holds them, for :meth:`fsdp_traffic`.
79 The non-expert parameters, the routed and the shared experts, each as
80 its count, the ranks the layer computes it over, the ranks that shard
81 it and the ranks that hold it, a shard or a copy.
82 """
83 non_exp, routed, shared = ctx.eval.num_p(ccfg, ctx)
84 return (
85 (non_exp, ccfg.t, ccfg.shard_p_os_non_exp_partial, ccfg.d * ccfg.cp * ccfg.t),
86 (routed / ccfg.ep, ccfg.t_exp, ccfg.shard_p_os_exp, max(1, int(ccfg.edp_group))),
87 (shared, ccfg.t_exp, ccfg.shard_p_os_exp_partial, ccfg.d * ccfg.cp * ccfg.t_exp),
88 )
90 @staticmethod
91 def fsdp_traffic(ccfg: CostModelConfig, parts: tuple, gathers: float, reduces: float) -> float:
92 """The bytes a rank moves in one micro-batch for a module FSDP shards.
94 Each ring collective over a part's n shards moves (n - 1) / n of the
95 part a rank. FSDP gathers a part's parameters *gathers* times a
96 micro-batch and reduce-scatters its gradient once, and the ranks
97 holding copies of a shard all-reduce it *reduces* times a micro-batch.
98 """
99 total = 0.0
100 for size, computed, kept, group in parts:
101 whole = size / max(1, computed)
102 shards = max(1.0, kept / max(1, computed))
103 copies = max(1.0, group / max(1, kept))
104 total += (shards - 1) / shards * whole * (gathers * ccfg.bytes_p + ccfg.bytes_grad)
105 total += 2 * (copies - 1) / copies * whole / shards * ccfg.bytes_grad * reduces
106 return total
108 @staticmethod
109 def fsdp_reduce_traffic(ccfg: CostModelConfig, parts: tuple, reduces: float) -> float:
110 """The all-reduce in :meth:`fsdp_traffic`: the bytes the copies of a shard move reducing it.
112 HSDP's replicas all-reduce each gradient shard once a step, and a run
113 whose FSDP shards nothing all-reduces every gradient; the rest of
114 :meth:`fsdp_traffic` is FSDP's own gathers and reduce-scatters.
115 """
116 total = 0.0
117 for size, computed, kept, group in parts:
118 whole = size / max(1, computed)
119 shards = max(1.0, kept / max(1, computed))
120 copies = max(1.0, group / max(1, kept))
121 total += 2 * (copies - 1) / copies * whole / shards * ccfg.bytes_grad * reduces
122 return total
124 @staticmethod
125 def tp_comm_non_exp(ccfg: CostModelConfig, ctx: Context, mb: int) -> float:
126 """TP comm for non-expert parameters"""
127 rec_layer = ctx.current_node == LayerType.SEL_REC_LAYER
128 tp_comm_non_exp = 0.25 * ccfg.n_gather
129 tp_comm_non_exp *= ccfg.s * ccfg.b * ccfg.h * mb
130 if ccfg.n_exp > 1:
131 tp_comm_non_exp = (
132 0.25
133 * ccfg.n_gather
134 * ccfg.h
135 * ccfg.h
136 * ccfg.bytes_compute
137 * ccfg.n_attMM
138 )
139 res = (
140 EvalUtils.rec_coeff(rec_layer, EvalUtils.switch(ccfg, ctx, "gather"))
141 * ccfg.comm_t
142 * tp_comm_non_exp
143 / ccfg.cp
144 )
145 return res
147 @staticmethod
148 def tp_comm_exp(ccfg: CostModelConfig, ctx: Context, mb: int) -> float:
149 """TP comm for expert parameters"""
150 rec_layer = ctx.current_node == LayerType.SEL_REC_LAYER
151 tp_comm_exp = 0.25 * ccfg.n_gather
152 tp_comm_exp *= ccfg.s * ccfg.b * ccfg.hff * mb
153 if ccfg.n_exp > 1:
154 # Routed and shared experts, each hff_exp wide
155 routed_comm = ccfg.n_exp / ccfg.ep * ccfg.hff_exp
156 shared_comm = ccfg.n_shared_exp * ccfg.hff_exp
157 tp_comm_exp = (
158 0.25
159 * ccfg.n_gather
160 * ccfg.h
161 * ccfg.bytes_compute
162 * ccfg.n_ffMM
163 * (routed_comm + shared_comm)
164 )
165 res = (
166 EvalUtils.rec_coeff(rec_layer, EvalUtils.switch(ccfg, ctx, "gather"))
167 * ccfg.comm_t
168 * tp_comm_exp
169 / ccfg.cp
170 )
171 return res
173 @staticmethod
174 def tp_comm_layer(ccfg: CostModelConfig, ctx: Context, mb: int) -> float:
175 """TP comm sum"""
176 non_exp = EvalLayerComm.tp_comm_non_exp(ccfg, ctx, mb)
177 exp = EvalLayerComm.tp_comm_exp(ccfg, ctx, mb)
178 return non_exp + exp
180 @staticmethod
181 def cp_comm_non_exp(ccfg: CostModelConfig, ctx: Context) -> float:
182 """CP comm for non-expert parameters"""
183 rec_layer = ctx.current_node == LayerType.SEL_REC_LAYER
184 rec_factor = EvalUtils.rec_coeff(rec_layer, EvalUtils.switch(ccfg, ctx, "gather")) * int(
185 ccfg.p == 1
186 ) # [HYPOTHESIS]
187 # hybird_cp is a known typo for hybrid_cp kept for backward compat
188 if ccfg.cp_algo in ["colossalai_cp", "hybrid_cp", "hybird_cp"]:
189 # FW Ring P2P + BW Ring P2P
190 # KV transfers, can be recomputed
191 return (
192 ccfg.comm_cp
193 * 2
194 * ccfg.s
195 * ccfg.b
196 * ((2 * 0.5 * rec_factor + 0.5) * ccfg.n_attMM * ccfg.h)
197 / (ccfg.t)
198 )
199 if ccfg.cp_algo == "ulysses_cp":
200 return (
201 ccfg.comm_cp
202 * 2
203 * ccfg.s
204 * ccfg.b
205 * ((0.5 * rec_factor + 0.5) * ccfg.n_attMM * ccfg.h)
206 / (ccfg.t)
207 )
208 return 0
210 @staticmethod
211 def cp_comm_exp(ccfg: CostModelConfig, _: Context) -> float:
212 """CP comm for expert parameters"""
213 # hybird_cp is a known typo for hybrid_cp kept for backward compat
214 if ccfg.cp_algo in ["colossalai_cp", "hybrid_cp", "hybird_cp", "ulysses_cp"]:
215 # FW Ring P2P + BW Ring P2P
216 # or FW + BW All2Alls
217 res = ccfg.comm_cp * 2 * ccfg.s * ccfg.b * ccfg.n_ffMM * ccfg.hff
218 return res / ccfg.t
219 return 0
221 @staticmethod
222 def cp_comm_layer(ccfg: CostModelConfig, ctx: Context) -> float:
223 """CP comm sum"""
224 non_exp = EvalLayerComm.cp_comm_non_exp(ccfg, ctx)
225 exp = EvalLayerComm.cp_comm_exp(ccfg, ctx)
226 return non_exp + exp
228 @staticmethod
229 def ep_comm_layer_balanced(
230 ccfg: CostModelConfig, ctx: Context, mb: int # pylint: disable=unused-argument
231 ) -> float:
232 """EP comm for balanced token distribution (byte volume).
234 Uses (ep-1)/ep correction: only (ep-1)/ep fraction of local tokens
235 actually cross rank boundaries in an all-to-all dispatch/combine pair.
236 Result is in bytes (like TP activation comm), unlike CP/DP which are
237 in element counts (parameter comm).
238 """
239 del ctx
240 if ccfg.ep <= 1 or ccfg.comm_ep == 0:
241 return 0
242 t_local = mb * ccfg.n_chosen_exp * ccfg.s * ccfg.b / ccfg.cp
243 t_cross = t_local * (ccfg.ep - 1) / ccfg.ep
244 return t_cross * ccfg.h * ccfg.bytes_compute * 2 * ccfg.comm_ep
246 @staticmethod
247 def ep_comm_layer_imbalanced(
248 ccfg: CostModelConfig, ctx: Context, mb: int
249 ) -> float:
250 """EP comm for imbalanced (skewed) token distribution (byte volume).
252 Uses max(rank_tokens) to bound communication volume.
253 Normalized with (ep-1)/ep cross-rank factor and mb scaling,
254 so it reduces to balanced when token distribution is uniform.
255 Falls back to balanced when tokens_per_expert is empty
256 or n_exp not divisible by ep.
258 tokens_per_expert: global per-expert token count per microbatch
259 (all EP ranks combined, before all-to-all dispatch; None = balanced).
260 Under uniform distribution, each rank's share equals
261 n_chosen_exp * s * b / (cp * t), matching t_local in the balanced formula.
263 Result is in bytes (like TP activation comm), unlike CP/DP which are
264 in element counts (parameter comm).
265 """
266 if ccfg.ep <= 1 or ccfg.comm_ep == 0:
267 return 0
268 tokens = ccfg.tokens_per_expert
269 if not tokens:
270 return EvalLayerComm.ep_comm_layer_balanced(ccfg, ctx, mb)
271 if ccfg.n_exp % ccfg.ep != 0:
272 logger.warning(
273 "n_exp=%d not divisible by ep=%d, falling back to balanced",
274 ccfg.n_exp,
275 ccfg.ep,
276 )
277 return EvalLayerComm.ep_comm_layer_balanced(ccfg, ctx, mb)
278 experts_per_rank = ccfg.n_exp // ccfg.ep
279 rank_tokens = []
280 for r in range(ccfg.ep):
281 rank_sum = sum(
282 tokens[r * experts_per_rank + i] for i in range(experts_per_rank)
283 )
284 rank_tokens.append(rank_sum)
285 max_inbound = max(rank_tokens)
286 # max_inbound: per-rank inbound tokens for one microbatch
287 # multiply by mb for the full pipeline stage, by (ep-1)/ep for cross-rank fraction
288 t_cross = max_inbound * mb * (ccfg.ep - 1) / ccfg.ep
289 return t_cross * ccfg.h * ccfg.bytes_compute * 2 * ccfg.comm_ep
291 @staticmethod
292 def ep_comm_layer(ccfg: CostModelConfig, ctx: Context, mb: int) -> float:
293 """EP comm dispatcher: balanced or imbalanced based on tokens_per_expert."""
294 if ccfg.ep <= 1 or ccfg.comm_ep == 0:
295 return 0
296 if ccfg.tokens_per_expert is not None:
297 return EvalLayerComm.ep_comm_layer_imbalanced(ccfg, ctx, mb)
298 return EvalLayerComm.ep_comm_layer_balanced(ccfg, ctx, mb)
300 @staticmethod
301 def cp_comm_buffer(ccfg: CostModelConfig, ctx: Context) -> float:
302 """Estimate CP communication buffer memory per layer.
304 Ring CP (colossalai_cp / hybrid_cp):
305 Intra-node: all-gather among device_per_node ranks → buffer for
306 (intra_ranks - 1) extra KV chunks.
307 Cross-node: intra-node all-gather result stays resident, plus 1
308 full-node KV chunk as ring receive buffer → peak is
309 (2 * intra_ranks - 1) extra chunks.
311 Ulysses CP:
312 All2All rearranges heads across CP ranks. Peak buffer is
313 1 send chunk + 1 receive chunk = 2 chunks, where each chunk
314 is the per-rank activation slice being exchanged.
315 """
316 del ctx
317 if ccfg.cp <= 1:
318 return 0.0
320 s, b, cp = ccfg.s, ccfg.b, ccfg.cp
321 fp16_bytes = 2
322 kv_bytes = fp16_bytes * 2
324 cp_algo = _resolve_cp_algo(ccfg)
326 if cp_algo == CPAlgo.ULYSSES_CP:
327 kv_dim = compute_kv_dim(ccfg)
328 chunk = s * b * (kv_dim / cp) * kv_bytes
329 intra_ranks = min(int(cp), int(ccfg.device_per_node))
330 if cp <= ccfg.device_per_node:
331 extra_chunks = intra_ranks - 1
332 else:
333 extra_chunks = 2 * intra_ranks - 1
334 return extra_chunks * chunk
336 kv_dim = compute_kv_dim(ccfg)
337 chunk = (s / cp) * b * kv_dim * kv_bytes
339 intra_ranks = min(int(cp), int(ccfg.device_per_node))
341 if cp <= ccfg.device_per_node:
342 extra_chunks = intra_ranks - 1
343 else:
344 extra_chunks = 2 * intra_ranks - 1
346 return extra_chunks * chunk