« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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"""EP constraint validators for auto-parallel strategy search.
17Four constraints that any valid EP strategy must satisfy:
18 C1: n_experts % ep_degree == 0 (expert divisibility)
19 C2: hff_exp % t_exp == 0 (expert hidden dim divisibility, t_exp = etp or tp)
20 C3: dp*tp*pp*cp <= total_devices (device limit, EP borrows from DP)
21 C4: EP+PP stage expert memory <= device_capacity (stage feasibility)
22"""
23from __future__ import annotations
24from dataclasses import dataclass
25from inspect import Parameter, Signature
26from typing import Any
29@dataclass
30class ConstraintResult:
31 """Structured result for a constraint check."""
33 name: str
34 passed: bool
35 message: str
37 def __bool__(self):
38 return self.passed
41@dataclass(frozen=True)
42class StageExperts:
43 """Expert weights hosted by one pipeline stage.
45 Attributes:
46 n_moe_layers: Number of MoE layers in this PP stage.
47 n_exp: Number of experts per MoE layer.
48 h: Model hidden dimension.
49 hff_exp: Expert FFN hidden dimension (per expert).
50 bytes_p: Bytes per parameter.
51 n_ff_mm: Number of feedforward linear layers per expert
52 (SwiGLU: 3, standard MLP: 2).
53 """
55 n_moe_layers: int
56 n_exp: int
57 h: int
58 hff_exp: int
59 bytes_p: int
60 n_ff_mm: int = 3
63_LEGACY_STAGE_SIGNATURE = Signature([
64 Parameter(name, Parameter.POSITIONAL_OR_KEYWORD)
65 for name in ("n_moe_layers", "n_exp", "ep", "dp", "h", "hff_exp", "bytes_p", "device_capacity_gb")
66] + [
67 Parameter("zero_level", Parameter.POSITIONAL_OR_KEYWORD, default=2),
68 Parameter("t_exp", Parameter.POSITIONAL_OR_KEYWORD, default=1),
69 Parameter("n_ffMM", Parameter.POSITIONAL_OR_KEYWORD, default=3),
70])
73class EpConstraints:
74 """EP-specific constraint checks for auto-parallel strategy search."""
76 @staticmethod
77 def check_ep_divisibility(n_exp: int, ep: int) -> ConstraintResult:
78 """C1: n_experts must be divisible by ep_degree."""
79 if ep < 1:
80 return ConstraintResult(
81 "ep_divisibility", False,
82 f"ep={ep} must be >= 1")
83 if n_exp < 1:
84 return ConstraintResult(
85 "ep_divisibility", False,
86 f"n_exp={n_exp} must be >= 1")
87 if n_exp % ep != 0:
88 return ConstraintResult(
89 "ep_divisibility", False,
90 f"n_exp={n_exp} not divisible by ep={ep}, "
91 f"remainder={n_exp % ep}")
92 return ConstraintResult(
93 "ep_divisibility", True,
94 f"n_exp={n_exp}/ep={ep} = {n_exp // ep} experts/rank")
96 @staticmethod
97 def check_expert_hidden_divisibility(
98 hff_exp: int, t_exp: int
99 ) -> ConstraintResult:
100 """C2: hff_exp must be divisible by expert TP degree.
102 Args:
103 hff_exp: Expert FFN hidden dimension.
104 t_exp: Expert TP degree (= etp if etp > 1, else tp).
105 etp and tp are alternatives (not multiplicative).
106 """
107 if t_exp < 1:
108 return ConstraintResult(
109 "expert_hidden_divisibility", False,
110 f"t_exp={t_exp} must be >= 1")
111 if hff_exp <= 0:
112 return ConstraintResult(
113 "expert_hidden_divisibility", True,
114 f"hff_exp={hff_exp} (dense FFN, no expert TP needed)")
115 if hff_exp % t_exp != 0:
116 return ConstraintResult(
117 "expert_hidden_divisibility", False,
118 f"hff_exp={hff_exp} not divisible by t_exp={t_exp}, "
119 f"remainder={hff_exp % t_exp}")
120 return ConstraintResult(
121 "expert_hidden_divisibility", True,
122 f"hff_exp={hff_exp}/t_exp={t_exp} = {hff_exp // t_exp}")
124 @staticmethod
125 def check_device_limit(
126 dp: int, tp: int, pp: int, cp: int, total_devices: int
127 ) -> ConstraintResult:
128 """C3: dp*tp*pp*cp must not exceed total_devices.
130 EP borrows devices from DP and does not occupy independent
131 device slots, so it is not included in the device count.
132 """
133 required = dp * tp * pp * cp
134 if required > total_devices:
135 return ConstraintResult(
136 "device_limit", False,
137 f"dp={dp}*tp={tp}*pp={pp}*cp={cp}="
138 f"{required} > total_devices={total_devices}")
139 return ConstraintResult(
140 "device_limit", True,
141 f"{required}/{total_devices} devices used")
143 @staticmethod
144 def check_ep_pp_stage_feasibility(*args: Any, **kwargs: Any) -> ConstraintResult:
145 """Check stage memory using StageExperts or the legacy numeric arguments.
147 Args:
148 *args: A StageExperts object followed by ep, dp and capacity, or the
149 legacy order n_moe_layers, n_exp, ep, dp, h, hff_exp, bytes_p,
150 device_capacity_gb, zero_level, t_exp, n_ffMM.
151 **kwargs: Named fields in the selected form. Legacy n_ffMM remains
152 accepted and maps to StageExperts.n_ff_mm.
154 Returns:
155 Whether the stage expert memory fits the device capacity.
156 """
157 if (args and isinstance(args[0], StageExperts)) or "experts" in kwargs:
158 return EpConstraints._check_stage_experts(*args, **kwargs)
159 bound = _LEGACY_STAGE_SIGNATURE.bind(*args, **kwargs)
160 bound.apply_defaults()
161 values = bound.arguments
162 experts = StageExperts(
163 **{key: values.pop(key) for key in ("n_moe_layers", "n_exp", "h", "hff_exp", "bytes_p")},
164 n_ff_mm=values.pop("n_ffMM"),
165 )
166 return EpConstraints._check_stage_experts(experts, **values)
168 @staticmethod
169 def _check_stage_experts(
170 experts: StageExperts,
171 ep: int,
172 dp: int,
173 device_capacity_gb: float,
174 zero_level: int = 2,
175 t_exp: int = 1,
176 ) -> ConstraintResult:
177 """C4: EP+PP stage expert memory must fit within device capacity.
179 Args:
180 experts: Expert weights hosted by this PP stage.
181 ep: Expert parallelism degree.
182 dp: Data parallelism degree (for ZeRO sharding).
183 device_capacity_gb: Device memory capacity in GB.
184 zero_level: ZeRO optimizer sharding level (2 or 3).
185 t_exp: Expert TP degree (= etp if etp > 1, else tp).
186 """
187 if experts.n_moe_layers <= 0:
188 return ConstraintResult(
189 "ep_pp_stage_feasibility", True,
190 "Dense-only stage, no expert memory")
191 if ep < 1:
192 return ConstraintResult(
193 "ep_pp_stage_feasibility", False,
194 f"ep={ep} must be >= 1")
195 experts_per_rank = experts.n_exp / ep
196 params_per_expert = experts.n_ff_mm * experts.h * experts.hff_exp / max(t_exp, 1)
197 param_bytes = (
198 experts.n_moe_layers * experts_per_rank * params_per_expert * experts.bytes_p
199 )
200 os_bytes = param_bytes * 2
201 grad_bytes = param_bytes
202 if zero_level >= 3:
203 total = (param_bytes + os_bytes + grad_bytes) / max(dp, 1)
204 elif zero_level == 2:
205 total = param_bytes + os_bytes / max(dp, 1) + grad_bytes / max(dp, 1)
206 else:
207 total = param_bytes + os_bytes + grad_bytes
208 total_gb = total / 1e9
209 if total_gb > device_capacity_gb:
210 return ConstraintResult(
211 "ep_pp_stage_feasibility", False,
212 f"Stage expert mem={total_gb:.1f}GB > capacity="
213 f"{device_capacity_gb:.1f}GB "
214 f"(exp/rank={experts_per_rank:.0f}, ZeRO-{zero_level})")
215 return ConstraintResult(
216 "ep_pp_stage_feasibility", True,
217 f"Stage expert mem={total_gb:.1f}GB <= capacity="
218 f"{device_capacity_gb:.1f}GB")
220 @classmethod
221 def validate_all(cls, ccfg, total_devices: int,
222 device_capacity_gb: float) -> list:
223 """Run all EP constraint checks, return list of ConstraintResult."""
224 t_exp = max(ccfg.etp, 1) if ccfg.etp > 1 else max(ccfg.t, 1)
225 results = [
226 cls.check_ep_divisibility(ccfg.n_exp, ccfg.ep),
227 cls.check_expert_hidden_divisibility(ccfg.hff_exp, t_exp),
228 cls.check_device_limit(
229 ccfg.d, ccfg.t, ccfg.p, ccfg.cp, total_devices),
230 ]
231 if ccfg.p > 1 and ccfg.n_exp > 1:
232 n_moe_per_stage = getattr(ccfg, 'n_lay', 1) // ccfg.p
233 experts = StageExperts(
234 n_moe_layers=n_moe_per_stage,
235 n_exp=ccfg.n_exp,
236 h=ccfg.h,
237 hff_exp=ccfg.hff_exp,
238 bytes_p=ccfg.bytes_p,
239 n_ff_mm=max(getattr(ccfg, 'n_ffMM', 3), 1),
240 )
241 results.append(cls.check_ep_pp_stage_feasibility(
242 experts,
243 ep=ccfg.ep, dp=ccfg.d,
244 device_capacity_gb=device_capacity_gb,
245 zero_level=int(ccfg.comm_d_exp), t_exp=t_exp,
246 ))
247 return results