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