« 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"""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 tp_comm_non_exp(ccfg: CostModelConfig, ctx: Context, mb: int) -> float:
77 """TP comm for non-expert parameters"""
78 rec_layer = ctx.current_node == LayerType.SEL_REC_LAYER
79 tp_comm_non_exp = 0.25 * ccfg.n_gather
80 tp_comm_non_exp *= ccfg.s * ccfg.b * ccfg.h * mb
81 if ccfg.n_exp > 1:
82 tp_comm_non_exp = (
83 0.25
84 * ccfg.n_gather
85 * ccfg.h
86 * ccfg.h
87 * ccfg.bytes_compute
88 * ccfg.n_attMM
89 )
90 res = (
91 EvalUtils.rec_coeff(rec_layer, ccfg.rec_op.gather)
92 * ccfg.comm_t
93 * tp_comm_non_exp
94 / ccfg.cp
95 )
96 return res
98 @staticmethod
99 def tp_comm_exp(ccfg: CostModelConfig, ctx: Context, mb: int) -> float:
100 """TP comm for expert parameters"""
101 rec_layer = ctx.current_node == LayerType.SEL_REC_LAYER
102 tp_comm_exp = 0.25 * ccfg.n_gather
103 tp_comm_exp *= ccfg.s * ccfg.b * ccfg.hff * mb
104 if ccfg.n_exp > 1:
105 # Routed experts use hff_exp, shared experts use hff
106 routed_comm = ccfg.n_exp / ccfg.ep * ccfg.hff_exp
107 shared_comm = ccfg.n_shared_exp * ccfg.hff
108 tp_comm_exp = (
109 0.25
110 * ccfg.n_gather
111 * ccfg.h
112 * ccfg.bytes_compute
113 * ccfg.n_ffMM
114 * (routed_comm + shared_comm)
115 )
116 res = (
117 EvalUtils.rec_coeff(rec_layer, ccfg.rec_op.gather)
118 * ccfg.comm_t
119 * tp_comm_exp
120 / ccfg.cp
121 )
122 return res
124 @staticmethod
125 def tp_comm_layer(ccfg: CostModelConfig, ctx: Context, mb: int) -> float:
126 """TP comm sum"""
127 non_exp = EvalLayerComm.tp_comm_non_exp(ccfg, ctx, mb)
128 exp = EvalLayerComm.tp_comm_exp(ccfg, ctx, mb)
129 return non_exp + exp
131 @staticmethod
132 def cp_comm_non_exp(ccfg: CostModelConfig, ctx: Context) -> float:
133 """CP comm for non-expert parameters"""
134 rec_layer = ctx.current_node == LayerType.SEL_REC_LAYER
135 rec_factor = EvalUtils.rec_coeff(rec_layer, ccfg.rec_op.gather) * int(
136 ccfg.p == 1
137 ) # [HYPOTHESIS]
138 # hybird_cp is a known typo for hybrid_cp kept for backward compat
139 if ccfg.cp_algo in ["colossalai_cp", "hybrid_cp", "hybird_cp"]:
140 # FW Ring P2P + BW Ring P2P
141 # KV transfers, can be recomputed
142 return (
143 ccfg.comm_cp
144 * 2
145 * ccfg.s
146 * ccfg.b
147 * ((2 * 0.5 * rec_factor + 0.5) * ccfg.n_attMM * ccfg.h)
148 / (ccfg.t)
149 )
150 if ccfg.cp_algo == "ulysses_cp":
151 return (
152 ccfg.comm_cp
153 * 2
154 * ccfg.s
155 * ccfg.b
156 * ((0.5 * rec_factor + 0.5) * ccfg.n_attMM * ccfg.h)
157 / (ccfg.t)
158 )
159 return 0
161 @staticmethod
162 def cp_comm_exp(ccfg: CostModelConfig, _: Context) -> float:
163 """CP comm for expert parameters"""
164 # hybird_cp is a known typo for hybrid_cp kept for backward compat
165 if ccfg.cp_algo in ["colossalai_cp", "hybrid_cp", "hybird_cp", "ulysses_cp"]:
166 # FW Ring P2P + BW Ring P2P
167 # or FW + BW All2Alls
168 res = ccfg.comm_cp * 2 * ccfg.s * ccfg.b * ccfg.n_ffMM * ccfg.hff
169 return res / ccfg.t
170 return 0
172 @staticmethod
173 def cp_comm_layer(ccfg: CostModelConfig, ctx: Context) -> float:
174 """CP comm sum"""
175 non_exp = EvalLayerComm.cp_comm_non_exp(ccfg, ctx)
176 exp = EvalLayerComm.cp_comm_exp(ccfg, ctx)
177 return non_exp + exp
179 @staticmethod
180 def ep_comm_layer_balanced(
181 ccfg: CostModelConfig, ctx: Context, mb: int # pylint: disable=unused-argument
182 ) -> float:
183 """EP comm for balanced token distribution (byte volume).
185 Uses (ep-1)/ep correction: only (ep-1)/ep fraction of local tokens
186 actually cross rank boundaries in an all-to-all dispatch/combine pair.
187 Result is in bytes (like TP activation comm), unlike CP/DP which are
188 in element counts (parameter comm).
189 """
190 del ctx
191 if ccfg.ep <= 1 or ccfg.comm_ep == 0:
192 return 0
193 t_local = mb * ccfg.n_chosen_exp * ccfg.s * ccfg.b / ccfg.cp
194 t_cross = t_local * (ccfg.ep - 1) / ccfg.ep
195 return t_cross * ccfg.h * ccfg.bytes_compute * 2 * ccfg.comm_ep
197 @staticmethod
198 def ep_comm_layer_imbalanced(
199 ccfg: CostModelConfig, ctx: Context, mb: int
200 ) -> float:
201 """EP comm for imbalanced (skewed) token distribution (byte volume).
203 Uses max(rank_tokens) to bound communication volume.
204 Normalized with (ep-1)/ep cross-rank factor and mb scaling,
205 so it reduces to balanced when token distribution is uniform.
206 Falls back to balanced when tokens_per_expert is empty
207 or n_exp not divisible by ep.
209 tokens_per_expert: global per-expert token count per microbatch
210 (all EP ranks combined, before all-to-all dispatch; None = balanced).
211 Under uniform distribution, each rank's share equals
212 n_chosen_exp * s * b / (cp * t), matching t_local in the balanced formula.
214 Result is in bytes (like TP activation comm), unlike CP/DP which are
215 in element counts (parameter comm).
216 """
217 if ccfg.ep <= 1 or ccfg.comm_ep == 0:
218 return 0
219 tokens = ccfg.tokens_per_expert
220 if not tokens:
221 return EvalLayerComm.ep_comm_layer_balanced(ccfg, ctx, mb)
222 if ccfg.n_exp % ccfg.ep != 0:
223 logger.warning(
224 "n_exp=%d not divisible by ep=%d, falling back to balanced",
225 ccfg.n_exp,
226 ccfg.ep,
227 )
228 return EvalLayerComm.ep_comm_layer_balanced(ccfg, ctx, mb)
229 experts_per_rank = ccfg.n_exp // ccfg.ep
230 rank_tokens = []
231 for r in range(ccfg.ep):
232 rank_sum = sum(
233 tokens[r * experts_per_rank + i] for i in range(experts_per_rank)
234 )
235 rank_tokens.append(rank_sum)
236 max_inbound = max(rank_tokens)
237 # max_inbound: per-rank inbound tokens for one microbatch
238 # multiply by mb for the full pipeline stage, by (ep-1)/ep for cross-rank fraction
239 t_cross = max_inbound * mb * (ccfg.ep - 1) / ccfg.ep
240 return t_cross * ccfg.h * ccfg.bytes_compute * 2 * ccfg.comm_ep
242 @staticmethod
243 def ep_comm_layer(ccfg: CostModelConfig, ctx: Context, mb: int) -> float:
244 """EP comm dispatcher: balanced or imbalanced based on tokens_per_expert."""
245 if ccfg.ep <= 1 or ccfg.comm_ep == 0:
246 return 0
247 if ccfg.tokens_per_expert is not None:
248 return EvalLayerComm.ep_comm_layer_imbalanced(ccfg, ctx, mb)
249 return EvalLayerComm.ep_comm_layer_balanced(ccfg, ctx, mb)
251 @staticmethod
252 def cp_comm_buffer(ccfg: CostModelConfig, ctx: Context) -> float:
253 """Estimate CP communication buffer memory per layer.
255 Ring CP (colossalai_cp / hybrid_cp):
256 Intra-node: all-gather among device_per_node ranks → buffer for
257 (intra_ranks - 1) extra KV chunks.
258 Cross-node: intra-node all-gather result stays resident, plus 1
259 full-node KV chunk as ring receive buffer → peak is
260 (2 * intra_ranks - 1) extra chunks.
262 Ulysses CP:
263 All2All rearranges heads across CP ranks. Peak buffer is
264 1 send chunk + 1 receive chunk = 2 chunks, where each chunk
265 is the per-rank activation slice being exchanged.
266 """
267 del ctx
268 if ccfg.cp <= 1:
269 return 0.0
271 s, b, cp = ccfg.s, ccfg.b, ccfg.cp
272 fp16_bytes = 2
273 kv_bytes = fp16_bytes * 2
275 cp_algo = _resolve_cp_algo(ccfg)
277 if cp_algo == CPAlgo.ULYSSES_CP:
278 kv_dim = compute_kv_dim(ccfg)
279 chunk = s * b * (kv_dim / cp) * kv_bytes
280 intra_ranks = min(int(cp), int(ccfg.device_per_node))
281 if cp <= ccfg.device_per_node:
282 extra_chunks = intra_ranks - 1
283 else:
284 extra_chunks = 2 * intra_ranks - 1
285 return extra_chunks * chunk
287 kv_dim = compute_kv_dim(ccfg)
288 chunk = (s / cp) * b * kv_dim * kv_bytes
290 intra_ranks = min(int(cp), int(ccfg.device_per_node))
292 if cp <= ccfg.device_per_node:
293 extra_chunks = intra_ranks - 1
294 else:
295 extra_chunks = 2 * intra_ranks - 1
297 return extra_chunks * chunk