« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +0800
1# Copyright 2025 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
22if TYPE_CHECKING:
23 from hyper_parallel.auto_parallel.sapp_nd.nd.common.cost_model_preprocess import CostModelConfig
24 from hyper_parallel.auto_parallel.sapp_nd.memory_estimation._context import Context
27class EvalLayerComm:
28 """Communication volume formulas class"""
30 @staticmethod
31 def dp_comm_non_exp(ccfg: CostModelConfig, ctx: Context) -> float:
32 """DP/OP comm for non-expert parameters"""
33 non_exp, _, _ = ctx.eval.num_p(ccfg, ctx)
34 dp_comm_non_exp = 0
35 # Non expert ZeRO LvL 2
36 if ccfg.comm_d_non_exp == 2:
37 dp_comm_non_exp += non_exp / (ccfg.cp * ccfg.t)
38 dp_comm_non_exp += non_exp / ccfg.t
39 # Non expert ZeRO LvL 3
40 if ccfg.comm_d_non_exp == 3:
41 dp_comm_non_exp += non_exp / ccfg.t
42 return dp_comm_non_exp
44 @staticmethod
45 def dp_comm_exp(ccfg: CostModelConfig, ctx: Context) -> float:
46 """DP/OP comm for expert parameters"""
47 _, routed, shared = ctx.eval.num_p(ccfg, ctx)
48 exp_param_size = routed + shared
49 if exp_param_size == 0:
50 return 0
51 dp_comm_exp = 0
52 # Expert ZeRO LvL 2
53 if ccfg.comm_d_exp == 2:
54 dp_comm_exp += exp_param_size / (ccfg.cp * ccfg.t_exp * ccfg.ep)
55 dp_comm_exp += exp_param_size / max(ccfg.ep, ccfg.t_exp)
56 # Expert ZeRO LvL 3
57 if ccfg.comm_d_exp == 3:
58 dp_comm_exp += exp_param_size / (ccfg.cp * ccfg.t_exp * ccfg.ep)
59 return dp_comm_exp
61 @staticmethod
62 def dp_comm_layer(ccfg: CostModelConfig, ctx: Context) -> float:
63 """DP/OP comm sum"""
64 non_exp = EvalLayerComm.dp_comm_non_exp(ccfg, ctx)
65 exp = EvalLayerComm.dp_comm_exp(ccfg, ctx)
66 return non_exp + exp
68 @staticmethod
69 def tp_comm_non_exp(ccfg: CostModelConfig, ctx: Context, mb: int) -> float:
70 """TP comm for non-expert parameters"""
71 rec_layer = ctx.current_node == LayerType.SEL_REC_LAYER
72 tp_comm_non_exp = 0.25 * ccfg.n_gather
73 tp_comm_non_exp *= ccfg.s * ccfg.b * ccfg.h * mb
74 if ccfg.n_exp > 1:
75 tp_comm_non_exp = (
76 0.25
77 * ccfg.n_gather
78 * ccfg.h
79 * ccfg.h
80 * ccfg.bytes_compute
81 * ccfg.n_attMM
82 )
83 res = (
84 EvalUtils.rec_coeff(rec_layer, ccfg.rec_op.gather)
85 * ccfg.comm_t
86 * tp_comm_non_exp
87 / ccfg.cp
88 )
89 return res
91 @staticmethod
92 def tp_comm_exp(ccfg: CostModelConfig, ctx: Context, mb: int) -> float:
93 """TP comm for expert parameters"""
94 rec_layer = ctx.current_node == LayerType.SEL_REC_LAYER
95 tp_comm_exp = 0.25 * ccfg.n_gather
96 tp_comm_exp *= ccfg.s * ccfg.b * ccfg.hff * mb
97 if ccfg.n_exp > 1:
98 # Routed experts use hff_exp, shared experts use hff
99 routed_comm = ccfg.n_exp / ccfg.ep * ccfg.hff_exp
100 shared_comm = ccfg.n_shared_exp * ccfg.hff
101 tp_comm_exp = (
102 0.25
103 * ccfg.n_gather
104 * ccfg.h
105 * ccfg.bytes_compute
106 * ccfg.n_ffMM
107 * (routed_comm + shared_comm)
108 )
109 res = (
110 EvalUtils.rec_coeff(rec_layer, ccfg.rec_op.gather)
111 * ccfg.comm_t
112 * tp_comm_exp
113 / ccfg.cp
114 )
115 return res
117 @staticmethod
118 def tp_comm_layer(ccfg: CostModelConfig, ctx: Context, mb: int) -> float:
119 """TP comm sum"""
120 non_exp = EvalLayerComm.tp_comm_non_exp(ccfg, ctx, mb)
121 exp = EvalLayerComm.tp_comm_exp(ccfg, ctx, mb)
122 return non_exp + exp
124 @staticmethod
125 def cp_comm_non_exp(ccfg: CostModelConfig, ctx: Context) -> float:
126 """CP comm for non-expert parameters"""
127 rec_layer = ctx.current_node == LayerType.SEL_REC_LAYER
128 rec_factor = EvalUtils.rec_coeff(rec_layer, ccfg.rec_op.gather) * int(
129 ccfg.p == 1
130 )
131 if ccfg.cp_algo in ["colossalai_cp", "hybird_cp"]:
132 return (
133 ccfg.comm_cp
134 * 2
135 * ccfg.s
136 * ccfg.b
137 * ((2 * 0.5 * rec_factor + 0.5) * ccfg.n_attMM * ccfg.h)
138 / (ccfg.t)
139 )
140 if ccfg.cp_algo == "ulysses_cp":
141 return (
142 ccfg.comm_cp
143 * 2
144 * ccfg.s
145 * ccfg.b
146 * ((0.5 * rec_factor + 0.5) * ccfg.n_attMM * ccfg.h)
147 / (ccfg.t)
148 )
149 return 0
151 @staticmethod
152 def cp_comm_exp(ccfg: CostModelConfig, _) -> float:
153 """CP comm for expert parameters"""
154 if ccfg.cp_algo in ["colossalai_cp", "hybird_cp", "ulysses_cp"]:
155 res = ccfg.comm_cp * 2 * ccfg.s * ccfg.b * ccfg.n_ffMM * ccfg.hff
156 return res / ccfg.t
157 return 0
159 @staticmethod
160 def cp_comm_layer(ccfg: CostModelConfig, ctx: Context) -> float:
161 """CP comm sum"""
162 non_exp = EvalLayerComm.cp_comm_non_exp(ccfg, ctx)
163 exp = EvalLayerComm.cp_comm_exp(ccfg, ctx)
164 return non_exp + exp
166 @staticmethod
167 def ep_comm_layer_balanced(
168 ccfg: CostModelConfig, ctx: Context, mb: int # pylint: disable=unused-argument
169 ) -> float:
170 """EP comm for balanced token distribution (byte volume).
172 Uses (ep-1)/ep correction: only (ep-1)/ep fraction of local tokens
173 actually cross rank boundaries in an all-to-all dispatch/combine pair.
174 Result is in bytes (like TP activation comm), unlike CP/DP which are
175 in element counts (parameter comm).
176 """
177 if ccfg.ep <= 1 or ccfg.comm_ep == 0:
178 return 0
179 t_local = mb * ccfg.n_chosen_exp * ccfg.s * ccfg.b / ccfg.cp
180 t_cross = t_local * (ccfg.ep - 1) / ccfg.ep
181 return t_cross * ccfg.h * ccfg.bytes_compute * 2 * ccfg.comm_ep
183 @staticmethod
184 def ep_comm_layer_imbalanced(
185 ccfg: CostModelConfig, ctx: Context, mb: int
186 ) -> float:
187 """EP comm for imbalanced (skewed) token distribution (byte volume).
189 Uses max(rank_tokens) to bound communication volume.
190 Normalized with (ep-1)/ep cross-rank factor and mb scaling,
191 so it reduces to balanced when token distribution is uniform.
192 Falls back to balanced when tokens_per_expert is empty
193 or n_exp not divisible by ep.
195 tokens_per_expert: global per-expert token count per microbatch
196 (all EP ranks combined, before all-to-all dispatch; None = balanced).
197 Under uniform distribution, each rank's share equals
198 n_chosen_exp * s * b / (cp * t), matching t_local in the balanced formula.
200 Result is in bytes (like TP activation comm), unlike CP/DP which are
201 in element counts (parameter comm).
202 """
203 if ccfg.ep <= 1 or ccfg.comm_ep == 0:
204 return 0
205 tokens = ccfg.tokens_per_expert
206 if not tokens:
207 return EvalLayerComm.ep_comm_layer_balanced(ccfg, ctx, mb)
208 if ccfg.n_exp % ccfg.ep != 0:
209 logger.warning(
210 "n_exp=%d not divisible by ep=%d, falling back to balanced",
211 ccfg.n_exp,
212 ccfg.ep,
213 )
214 return EvalLayerComm.ep_comm_layer_balanced(ccfg, ctx, mb)
215 experts_per_rank = ccfg.n_exp // ccfg.ep
216 rank_tokens = []
217 for r in range(ccfg.ep):
218 rank_sum = sum(
219 tokens[r * experts_per_rank + i] for i in range(experts_per_rank)
220 )
221 rank_tokens.append(rank_sum)
222 max_inbound = max(rank_tokens)
223 # max_inbound: per-rank inbound tokens for one microbatch
224 # multiply by mb for the full pipeline stage, by (ep-1)/ep for cross-rank fraction
225 t_cross = max_inbound * mb * (ccfg.ep - 1) / ccfg.ep
226 return t_cross * ccfg.h * ccfg.bytes_compute * 2 * ccfg.comm_ep
228 @staticmethod
229 def ep_comm_layer(ccfg: CostModelConfig, ctx: Context, mb: int) -> float:
230 """EP comm dispatcher: balanced or imbalanced based on tokens_per_expert."""
231 if ccfg.ep <= 1 or ccfg.comm_ep == 0:
232 return 0
233 if ccfg.tokens_per_expert is not None:
234 return EvalLayerComm.ep_comm_layer_imbalanced(ccfg, ctx, mb)
235 return EvalLayerComm.ep_comm_layer_balanced(ccfg, ctx, mb)