Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / expert_parallel / recipes.py: 87%
67 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07: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"""expert_parallel.recipes: built-in EP compute factories, organized as EXPLICIT archetypes
16(accuracy_fix_plan.md §3 E2).
18Since the explicit-injection rework the planner only SHARDS the expert
19parameters (``{EP: Shard(0)}`` + stacking metadata); the EP compute is never
20injected automatically — a factory must be referenced explicitly from YAML:
22.. code-block:: yaml
24 plan_overrides:
25 - match: "*.mlp.forward"
26 when: ep
27 replace:
28 kind: method
29 implementation: hyper_parallel.distributed.expert_parallel.recipes.qwen2moe_ep_compute_fn
31Archetype table (each factory is a COMPLETE semantic implementation —
32router / dispatch / experts / combine / shared expert / gate / merge are
33all inside; the framework never fills in any branch for you):
35=============================== ====================== =========================== =========================
36archetype key router shared expert representative models
37=============================== ====================== =========================== =========================
38routed_only_softmax_topk softmax top-k none legacy linear routers
39mixtral_topk_router TopKRouter tuple none Mixtral
40qwen2moe_shared_expert_gate TopKRouter tuple shared_expert + sigmoid gate Qwen2-MoE
41deepseekv3_sigmoid_group_shared sigmoid group routing shared_experts DeepSeek-V3 / GLM-4-MoE
42=============================== ====================== =========================== =========================
44Model-family archetypes live at their model-adapter homes (adjust doc §5.1):
45the Qwen3-MoE ``qwen3moe_topk_router`` archetype is
46``models/qwen3_moe/adapter/distributed/expert_parallel.qwen3moe_ep_compute_fn``.
48Three hard constraints on every archetype (accuracy_fix_plan.md §3.2):
501. **Explicitly keyed, zero runtime probing** — the user names the factory
51 in YAML; the implementation accesses ``module.experts`` /
52 ``module.shared_expert`` etc. directly. There are no getattr fallback
53 chains anywhere (a half-found branch is a SILENT numeric error — see
54 accuracy_problem.md 10.2).
552. **Apply-time interface assertion** — the factory runs ONCE at apply time
56 with the real module; it first asserts every attribute it will call and
57 fails the build with a teaching error (listing the module's actual
58 children) on mismatch. The assertion proves structural match, not
59 semantic correctness (that is vouched by numeric verification).
603. **Complete semantics** — nothing is left for the framework to guess.
62Not sure which archetype fits? ``describe_moe_module(module)`` (expert_parallel.experts)
63prints the module's children/param shapes; ``EP_ARCHETYPE_SUGGESTIONS``
64maps the planner's canonical arch name to a likely archetype (a SUGGESTION
65shown in error messages — selection stays explicit; a wrong suggestion is
66caught by constraint 2). Non-typical MoE: copy
67``examples/distributed/ep_factories.py`` and write your own factory with
68the same contract.
70Nested-boundary call contract (inside every compute_fn): calling a
71planned nested-boundary submodule (e.g. ``module.shared_expert(...)``)
72returns its out_dst layout with TP communication already sealed inside —
73the caller MUST NOT apply compensating collectives to the return value:
75.. code-block:: python
77 dist.all_reduce(shared, group=tp_group) # FORBIDDEN — the nested
78 # boundary already reduced
80Factory contract: decorated ``@local_compute`` (injection discipline —
81undecorated factories fail fast in the resolution chain). The mesh family
82``mesh`` / ``tp_mesh`` / ``cp_mesh`` / ``ep_mesh`` is MANDATORY context,
83ALL filled by the framework at apply time (None for inactive axes;
84``ep_mesh`` is the same object the expert parameters were sharded on).
85The factory must RETURN the compute fn
86``fn(module, *local_args, **local_kwargs) -> Tensor`` executed on local
87tensors inside the local-region skeleton.
88"""
90from typing import Any, Callable, Optional
92import torch
94from hyper_parallel.distributed.expert_parallel.routing import (
95 resolve_moe_router,
96)
97from hyper_parallel.distributed.expert_parallel.experts import (
98 bind_local_expert_forward,
99 ep_routed_dispatch,
100 ep_routed_experts_and_combine,
101 ep_routed_forward,
102 require_attrs,
103)
104from hyper_parallel.distributed.recipe_spec import (
105 local_compute,
106)
109# {canonical arch name (ShardingPlanner._get_architecture): archetype key} —
110# a SUGGESTION shown in preflight error messages and describe_moe_module
111# output, never an automatic selection (accuracy_fix_plan.md §3 E2).
112EP_ARCHETYPE_SUGGESTIONS = {
113 "qwen2moe": "qwen2moe_shared_expert_gate",
114 "qwen2_moe": "qwen2moe_shared_expert_gate",
115 "mixtral": "mixtral_topk_router",
116 "deepseekv3": "deepseekv3_sigmoid_group_shared",
117 "deepseek_v3": "deepseekv3_sigmoid_group_shared",
118 "glm4moe": "deepseekv3_sigmoid_group_shared",
119 "glm4_moe": "deepseekv3_sigmoid_group_shared",
120}
123def _require_ep_group(ep_mesh: Any, factory_name: str):
124 """Validate the expert mesh context and return the EP process group."""
125 if ep_mesh is None:
126 raise ValueError(
127 f"{factory_name} was built without an ep_mesh — this factory must "
128 "be injected on an EP-sharded MoE boundary (ep_size > 1 so the "
129 "framework-derived expert mesh exists). If you meant TP-only "
130 "expert sharding, do not use an EP compute factory."
131 )
132 return ep_mesh.get_group("ep")
135def _require_moe_interface(module: Any, expected_attrs, archetype_key: str) -> None:
136 """Apply-time structural assertion for an archetype (constraint 2)."""
137 require_attrs(module, *expected_attrs, owner=f"archetype '{archetype_key}'")
140def build_ep_compute(
141 module: Any,
142 ep_mesh: Any,
143 *,
144 router_fn: Callable,
145 archetype_key: str,
146 expected_attrs,
147 combine: Callable,
148 use_grouped_gemm: bool = False,
149 fix_router: bool = False,
150 capacity_factor: Optional[float] = None,
151 overlap_fn: Optional[Callable] = None,
152) -> Callable:
153 """Shared skeleton for archetype factories: validate context, assert the
154 interface, bind the local expert entry point, and close over the
155 archetype-specific ``combine`` (the branch composition — the ONLY part
156 that differs between archetypes).
158 ``combine(module, hidden_states, routed) -> Tensor`` receives the routed
159 branch output and returns the MoE block output.
161 When ``overlap_fn`` is given the routed branch is split into dispatch and
162 experts+combine, and ``overlap_fn(module, hidden_states)`` runs in between:
163 the token exchange is issued asynchronously there, so the overlapped work
164 (for DeepSeek-V3 style blocks the shared-expert MLP, which depends only on
165 ``hidden_states``) proceeds while it is in flight. ``combine`` then
166 receives that value as a fourth argument.
168 Public since M3 (adjust doc §5.4): model adapters (e.g.
169 ``models/qwen3_moe/adapter/distributed/expert_parallel.py``) compose
170 their EP archetype on top of this skeleton while the router adapters and
171 dispatch primitives stay in this generic layer.
172 ``fix_router`` is the archetype-independent benchmark switch: it swaps the chosen
173 router for its balanced variant -- the ``<name>_fixed`` registry entry when
174 ``router_fn`` names one, else
175 :func:`~hyper_parallel.distributed.expert_parallel.routing.make_balanced_router`
176 around the given callable -- so a family gets the uniform-load comparison mode
177 without carrying its own copy of the router. See that factory for what it does and
178 why its numbers must not be reported as an efficiency result.
179 """
180 router_fn = resolve_moe_router(router_fn, fix_router)
181 ep_group = _require_ep_group(ep_mesh, f"archetype '{archetype_key}'")
182 _require_moe_interface(module, expected_attrs, archetype_key)
183 if use_grouped_gemm:
184 bind_local_expert_forward(
185 module,
186 ep_mesh["ep"].size(),
187 use_grouped_gemm=True,
188 )
189 else:
190 # Keep the established call contract unchanged for every reference
191 # archetype and for Qwen3 when grouped GEMM is disabled.
192 bind_local_expert_forward(module, ep_mesh["ep"].size())
194 if capacity_factor is not None:
195 # Read back by ``experts._resolve_capacity_factor``; bounds each expert's routed slots
196 # so the busiest rank's buffers cannot blow up under a skewed real router.
197 module.experts.capacity_factor = float(capacity_factor)
199 if overlap_fn is None:
200 def compute_fn(module: Any, hidden_states: torch.Tensor) -> torch.Tensor:
201 """Run the routed branch and compose the MoE block output."""
202 routed = ep_routed_forward(
203 module, hidden_states, router_fn=router_fn, ep_group=ep_group)
204 return combine(module, hidden_states, routed)
205 else:
206 def compute_fn(module: Any, hidden_states: torch.Tensor) -> torch.Tensor:
207 """Dispatch, run ``overlap_fn`` against the in-flight exchange, then finish."""
208 state = ep_routed_dispatch(
209 module, hidden_states, router_fn=router_fn, ep_group=ep_group)
210 overlapped = overlap_fn(module, hidden_states)
211 routed = ep_routed_experts_and_combine(module, state, ep_group)
212 return combine(module, hidden_states, routed, overlapped)
214 return compute_fn
217@local_compute
218def routed_only_ep_compute_fn(
219 *,
220 module: Any,
221 mesh: Any,
222 tp_mesh: Any,
223 cp_mesh: Any,
224 ep_mesh: Any,
225 fix_router: bool = False,
226) -> Callable:
227 """Archetype ``routed_only_softmax_topk``: softmax top-k router, no
228 shared expert (Mixtral-style). Output = routed branch only.
230 Expected module interface: ``gate`` (router Linear), ``experts``.
232 ``fix_router`` selects the benchmark-only balanced variant of the router
233 (``routing.make_balanced_router``); correctness runs keep it ``False``.
234 """
235 del mesh, tp_mesh, cp_mesh
236 return build_ep_compute(
237 module,
238 ep_mesh,
239 router_fn=resolve_moe_router("default", fix_router),
240 archetype_key="routed_only_softmax_topk",
241 expected_attrs=["gate", "experts"],
242 combine=lambda module, hidden_states, routed: routed,
243 )
246@local_compute
247def mixtral_ep_compute_fn(
248 *,
249 module: Any,
250 mesh: Any,
251 tp_mesh: Any,
252 cp_mesh: Any,
253 ep_mesh: Any,
254 fix_router: bool = False,
255) -> Callable:
256 """Archetype ``mixtral_topk_router`` with training-time router jitter.
258 Transformers 5.12 Mixtral uses a TopKRouter returning
259 ``(logits, scores, indices)``. When configured, the original block also
260 perturbs hidden states before both routing and expert computation.
262 ``fix_router`` selects the benchmark-only balanced variant of the router
263 (``routing.make_balanced_router``); correctness runs keep it ``False``.
264 """
265 del mesh, tp_mesh, cp_mesh
266 router_fn = resolve_moe_router("mixtral", fix_router)
267 ep_group = _require_ep_group(ep_mesh, "archetype 'mixtral_topk_router'")
268 _require_moe_interface(module, ["gate", "experts"], "mixtral_topk_router")
269 bind_local_expert_forward(module, ep_mesh["ep"].size())
271 def compute_fn(module: Any, hidden_states: torch.Tensor) -> torch.Tensor:
272 """Apply router jitter when training, then run the routed branch."""
273 jitter_noise = float(getattr(module, "jitter_noise", 0.0))
274 if module.training and jitter_noise > 0.0:
275 noise = torch.empty_like(hidden_states).uniform_(
276 1.0 - jitter_noise,
277 1.0 + jitter_noise,
278 )
279 hidden_states = hidden_states * noise
280 return ep_routed_forward(
281 module,
282 hidden_states,
283 router_fn=router_fn,
284 ep_group=ep_group,
285 )
287 return compute_fn
290@local_compute
291def qwen2moe_ep_compute_fn(
292 *,
293 module: Any,
294 mesh: Any,
295 tp_mesh: Any,
296 cp_mesh: Any,
297 ep_mesh: Any,
298 fix_router: bool = False,
299) -> Callable:
300 """Archetype ``qwen2moe_shared_expert_gate``: TopKRouter tuple +
301 shared_expert + sigmoid scalar gate (Qwen2-MoE):
302 ``routed + sigmoid(shared_expert_gate(x)) * shared_expert(x)``.
304 Expected module interface: ``gate``, ``experts``, ``shared_expert``,
305 ``shared_expert_gate``.
307 ``module.shared_expert(...)`` is a planned nested TP boundary: its exit
308 already performs the RowWise Partial TP reduction (nested-boundary call
309 contract) — do NOT all-reduce its return value.
310 """
311 del mesh, tp_mesh, cp_mesh
313 def combine(
314 module: Any,
315 hidden_states: torch.Tensor,
316 routed: torch.Tensor,
317 ) -> torch.Tensor:
318 """Merge the routed branch with the gated shared-expert branch."""
319 shared = module.shared_expert(hidden_states) # nested boundary
320 gate = torch.sigmoid(module.shared_expert_gate(hidden_states))
321 return routed + gate * shared
323 return build_ep_compute(
324 module,
325 ep_mesh,
326 router_fn="qwen2moe",
327 fix_router=fix_router,
328 archetype_key="qwen2moe_shared_expert_gate",
329 expected_attrs=["gate", "experts", "shared_expert", "shared_expert_gate"],
330 combine=combine,
331 )
334@local_compute
335def deepseekv3_ep_compute_fn(
336 *,
337 module: Any,
338 mesh: Any,
339 tp_mesh: Any,
340 cp_mesh: Any,
341 ep_mesh: Any,
342 use_grouped_gemm: bool = False,
343 overlap_shared_expert: bool = False,
344 fix_router: bool = False,
345 capacity_factor: Optional[float] = None,
346) -> Callable:
347 """Archetype ``deepseekv3_sigmoid_group_shared``: sigmoid group-limited
348 routing (with e_score_correction_bias / routed_scaling_factor, already
349 folded into the router adapter's topk weights) + shared_experts
350 (DeepSeek-V3 / GLM-4-MoE): ``routed + shared_experts(x)``.
352 Expected module interface: ``gate``, ``experts``, ``shared_experts``.
354 ``module.shared_experts(...)`` is a planned nested TP boundary — its
355 exit already performs the TP reduction (nested-boundary call contract).
357 ``use_grouped_gemm=True`` runs the local experts through the packed
358 ``gate_up_proj`` grouped GEMM path (``npu_grouped_swiglu``) instead of
359 the eager per-expert loop; the default keeps the eager path.
361 ``overlap_shared_expert=True`` issues the token exchange asynchronously
362 and evaluates ``shared_experts`` while it is in flight — the shared branch
363 reads only ``hidden_states``, so it hides communication instead of
364 serializing behind the routed branch. The math is unchanged either way.
366 ``fix_router=True`` swaps the adapter for ``deepseekv3_fixed``, a
367 BENCHMARK-ONLY variant that hands every rank the same token count
368 (round-robin over the K slots by destination rank instead of the gate's
369 top-k) so that per-rank step time stops depending on how the data
370 happened to route. Loss / grad_norm / expert-load metrics are meaningless
371 in that mode and it must not be used to report an efficiency number —
372 see ``routing.make_balanced_router``. Correctness runs keep it off.
373 """
374 del mesh, tp_mesh, cp_mesh
376 def combine(
377 module: Any,
378 hidden_states: torch.Tensor,
379 routed: torch.Tensor,
380 overlapped: Optional[torch.Tensor] = None,
381 ) -> torch.Tensor:
382 """Merge the routed branch with the shared-experts branch."""
383 shared = module.shared_experts(hidden_states) if overlapped is None else overlapped
384 return routed + shared # nested boundary
386 def overlap_shared(module: Any, hidden_states: torch.Tensor) -> torch.Tensor:
387 """Shared-experts branch, evaluated while the dispatch exchange is in flight."""
388 return module.shared_experts(hidden_states) # nested boundary
390 return build_ep_compute(
391 module,
392 ep_mesh,
393 router_fn="deepseekv3",
394 fix_router=fix_router,
395 archetype_key="deepseekv3_sigmoid_group_shared",
396 expected_attrs=["gate", "experts", "shared_experts"],
397 combine=combine,
398 use_grouped_gemm=use_grouped_gemm,
399 capacity_factor=capacity_factor,
400 overlap_fn=overlap_shared if overlap_shared_expert else None,
401 )
404# {archetype key: factory} — introspection registry (the YAML _target_
405# points at a factory directly; this registry exists for
406# describe/preflight messaging).
407EP_ARCHETYPES = {
408 "routed_only_softmax_topk": routed_only_ep_compute_fn,
409 "mixtral_topk_router": mixtral_ep_compute_fn,
410 "qwen2moe_shared_expert_gate": qwen2moe_ep_compute_fn,
411 "deepseekv3_sigmoid_group_shared": deepseekv3_ep_compute_fn,
412}