« prev ^ index » next coverage.py v7.13.1, created at 2026-08-04 05:18 +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
27@dataclass
28class ConstraintResult:
29 """Structured result for a constraint check."""
31 name: str
32 passed: bool
33 message: str
35 def __bool__(self):
36 return self.passed
39class EpConstraints:
40 """EP-specific constraint checks for auto-parallel strategy search."""
42 @staticmethod
43 def check_ep_divisibility(n_exp: int, ep: int) -> ConstraintResult:
44 """C1: n_experts must be divisible by ep_degree."""
45 if ep < 1:
46 return ConstraintResult(
47 "ep_divisibility", False,
48 f"ep={ep} must be >= 1")
49 if n_exp < 1:
50 return ConstraintResult(
51 "ep_divisibility", False,
52 f"n_exp={n_exp} must be >= 1")
53 if n_exp % ep != 0:
54 return ConstraintResult(
55 "ep_divisibility", False,
56 f"n_exp={n_exp} not divisible by ep={ep}, "
57 f"remainder={n_exp % ep}")
58 return ConstraintResult(
59 "ep_divisibility", True,
60 f"n_exp={n_exp}/ep={ep} = {n_exp // ep} experts/rank")
62 @staticmethod
63 def check_expert_hidden_divisibility(
64 hff_exp: int, t_exp: int
65 ) -> ConstraintResult:
66 """C2: hff_exp must be divisible by expert TP degree.
68 Args:
69 hff_exp: Expert FFN hidden dimension.
70 t_exp: Expert TP degree (= etp if etp > 1, else tp).
71 etp and tp are alternatives (not multiplicative).
72 """
73 if t_exp < 1:
74 return ConstraintResult(
75 "expert_hidden_divisibility", False,
76 f"t_exp={t_exp} must be >= 1")
77 if hff_exp <= 0:
78 return ConstraintResult(
79 "expert_hidden_divisibility", True,
80 f"hff_exp={hff_exp} (dense FFN, no expert TP needed)")
81 if hff_exp % t_exp != 0:
82 return ConstraintResult(
83 "expert_hidden_divisibility", False,
84 f"hff_exp={hff_exp} not divisible by t_exp={t_exp}, "
85 f"remainder={hff_exp % t_exp}")
86 return ConstraintResult(
87 "expert_hidden_divisibility", True,
88 f"hff_exp={hff_exp}/t_exp={t_exp} = {hff_exp // t_exp}")
90 @staticmethod
91 def check_device_limit(
92 dp: int, tp: int, pp: int, cp: int, total_devices: int
93 ) -> ConstraintResult:
94 """C3: dp*tp*pp*cp must not exceed total_devices.
96 EP borrows devices from DP and does not occupy independent
97 device slots, so it is not included in the device count.
98 """
99 required = dp * tp * pp * cp
100 if required > total_devices:
101 return ConstraintResult(
102 "device_limit", False,
103 f"dp={dp}*tp={tp}*pp={pp}*cp={cp}="
104 f"{required} > total_devices={total_devices}")
105 return ConstraintResult(
106 "device_limit", True,
107 f"{required}/{total_devices} devices used")
109 @staticmethod
110 def check_ep_pp_stage_feasibility(
111 n_moe_layers: int,
112 n_exp: int,
113 ep: int,
114 dp: int,
115 h: int,
116 hff_exp: int,
117 bytes_p: int,
118 device_capacity_gb: float,
119 zero_level: int = 2,
120 t_exp: int = 1,
121 n_ffMM: int = 3, # pylint: disable=invalid-name
122 ) -> ConstraintResult:
123 """C4: EP+PP stage expert memory must fit within device capacity.
125 Args:
126 n_moe_layers: Number of MoE layers in this PP stage.
127 n_exp: Number of experts per MoE layer.
128 ep: Expert parallelism degree.
129 dp: Data parallelism degree (for ZeRO sharding).
130 h: Model hidden dimension.
131 hff_exp: Expert FFN hidden dimension (per expert).
132 bytes_p: Bytes per parameter.
133 device_capacity_gb: Device memory capacity in GB.
134 zero_level: ZeRO optimizer sharding level (2 or 3).
135 t_exp: Expert TP degree (= etp if etp > 1, else tp).
136 n_ffMM: Number of feedforward linear layers per expert
137 (SwiGLU: 3, standard MLP: 2).
138 """
139 if n_moe_layers <= 0:
140 return ConstraintResult(
141 "ep_pp_stage_feasibility", True,
142 "Dense-only stage, no expert memory")
143 if ep < 1:
144 return ConstraintResult(
145 "ep_pp_stage_feasibility", False,
146 f"ep={ep} must be >= 1")
147 experts_per_rank = n_exp / ep
148 params_per_expert = n_ffMM * h * hff_exp / max(t_exp, 1)
149 param_bytes = n_moe_layers * experts_per_rank * params_per_expert * bytes_p
150 os_bytes = param_bytes * 2
151 grad_bytes = param_bytes
152 if zero_level >= 3:
153 total = (param_bytes + os_bytes + grad_bytes) / max(dp, 1)
154 elif zero_level == 2:
155 total = param_bytes + os_bytes / max(dp, 1) + grad_bytes / max(dp, 1)
156 else:
157 total = param_bytes + os_bytes + grad_bytes
158 total_gb = total / 1e9
159 if total_gb > device_capacity_gb:
160 return ConstraintResult(
161 "ep_pp_stage_feasibility", False,
162 f"Stage expert mem={total_gb:.1f}GB > capacity="
163 f"{device_capacity_gb:.1f}GB "
164 f"(exp/rank={experts_per_rank:.0f}, ZeRO-{zero_level})")
165 return ConstraintResult(
166 "ep_pp_stage_feasibility", True,
167 f"Stage expert mem={total_gb:.1f}GB <= capacity="
168 f"{device_capacity_gb:.1f}GB")
170 @classmethod
171 def validate_all(cls, ccfg, total_devices: int,
172 device_capacity_gb: float) -> list:
173 """Run all EP constraint checks, return list of ConstraintResult."""
174 t_exp = max(ccfg.etp, 1) if ccfg.etp > 1 else max(ccfg.t, 1)
175 results = [
176 cls.check_ep_divisibility(ccfg.n_exp, ccfg.ep),
177 cls.check_expert_hidden_divisibility(ccfg.hff_exp, t_exp),
178 cls.check_device_limit(
179 ccfg.d, ccfg.t, ccfg.p, ccfg.cp, total_devices),
180 ]
181 if ccfg.p > 1 and ccfg.n_exp > 1:
182 n_moe_per_stage = getattr(ccfg, 'n_lay', 1) // ccfg.p
183 results.append(cls.check_ep_pp_stage_feasibility(
184 n_moe_layers=n_moe_per_stage,
185 n_exp=ccfg.n_exp,
186 ep=ccfg.ep, dp=ccfg.d, h=ccfg.h,
187 hff_exp=ccfg.hff_exp, bytes_p=ccfg.bytes_p,
188 device_capacity_gb=device_capacity_gb,
189 zero_level=int(ccfg.comm_d_exp), t_exp=t_exp,
190 n_ffMM=max(getattr(ccfg, 'n_ffMM', 3), 1),
191 ))
192 return results