Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / expert_parallel / expert_parallel.py: 86%
357 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-13 05:07 +0800
1# Copyright 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 Parallelism distributed strategies.
17Provides token permutation helpers and four parallel styles that compose with
18:class:`~hyper_parallel.core.expert_parallel.moe.GroupedExperts`:
20- :class:`BaseExpertParallel` — abstract base for EP strategies with
21 all-to-all token dispatch/combine.
22- :class:`ExpertParallel` — standard EP: each rank owns a shard of experts;
23 tokens are routed via differentiable all-to-all.
24- :class:`TensorParallel` — TP-only weight sharding for experts with no token
25 dispatch; for use when EP degree = 1.
26- :class:`ExpertTensorParallel` — combined EP + TP on a 2-D mesh ``[ep, tp]``;
27 weights are doubly sharded, dispatch uses the EP sub-mesh.
28"""
29__all__ = [
30 "AllToAllTokenDispatcher",
31 "DeredundencyTokenDispatcher",
32 "BaseExpertParallel",
33 "ExpertParallel",
34 "TensorParallel",
35 "ExpertTensorParallel",
36]
38from abc import ABC, abstractmethod
39from dataclasses import dataclass
40from typing import Any, List, Optional, Tuple, Union
42from hyper_parallel.core.dtensor.device_mesh import DeviceMesh
43from hyper_parallel.core.dtensor.dtensor import (
44 distribute_module,
45 distribute_tensor,
46 _distribute_module_iter_params,
47 _distribute_module_new_parameter,
48 _distribute_module_param_source,
49 _distribute_module_set_param,
50)
51from hyper_parallel.core.dtensor.placement_types import Shard
52from hyper_parallel.core.tensor_parallel.style import ParallelStyle
53from hyper_parallel.platform import AsyncHandle, get_platform
55platform = get_platform()
56Module = platform.Module
59# ---------------------------------------------------------------------------
60# Token permutation helpers
61# ---------------------------------------------------------------------------
63def _generate_permute_indices(
64 tokens_per_expert_group,
65 experts_per_rank: int,
66 num_ranks: int,
67):
68 """Generate permutation indices for rank-major → expert-major reordering.
70 After all-to-all, received tokens are laid out in rank-major order::
72 [rank0·expert0 tokens | rank0·expert1 tokens | ... |
73 rank1·expert0 tokens | rank1·expert1 tokens | ...]
75 Expert computation requires expert-major order::
77 [all tokens for local expert 0 | all tokens for local expert 1 | ...]
79 Args:
80 tokens_per_expert_group: 1-D integer tensor of shape
81 ``[num_ranks * experts_per_rank]``. Entry ``[r * E + e]`` is the
82 number of tokens received from rank ``r`` for local expert ``e``.
83 experts_per_rank: Number of experts owned by each rank.
84 num_ranks: EP degree (total number of ranks in the EP group).
86 Returns:
87 Tuple of:
89 - ``permuted_indices``: 1-D long tensor of length
90 ``total_received_tokens``. ``permuted_indices[i]`` is the source
91 position in the rank-major buffer for destination position ``i`` in
92 the expert-major buffer.
93 - ``num_tokens_per_expert``: 1-D integer tensor of length
94 ``experts_per_rank`` with the token count per local expert.
95 """
96 counts = tokens_per_expert_group # [num_ranks * experts_per_rank]
98 # num_tokens_per_expert[e] = Σ_r counts[r * E + e]
99 counts_2d = counts.view(num_ranks, experts_per_rank) # [R, E]
100 num_tokens_per_expert = counts_2d.sum(dim=0) # [E]
102 # ``total`` must be a host int because ``arange`` needs a scalar size.
103 # That single D2H drain is unavoidable. Everything else stays on
104 # device — no per-block ``.item()`` in a loop.
105 total = int(num_tokens_per_expert.sum())
106 if total == 0:
107 return counts.new_zeros(0, dtype=counts.dtype), num_tokens_per_expert
109 # ---- Vectorized expert-major permutation, no host stalls -----------
110 # Source offsets in the rank-major receive buffer for each (r, e) block.
111 src_offsets_rm = counts.cumsum(0) - counts # [R*E], starts of each block
112 # Reorder src offsets to expert-major iteration order: block (e, r).
113 src_offsets_em = (
114 src_offsets_rm.view(num_ranks, experts_per_rank).T.contiguous().view(-1)
115 ) # [E*R]
116 # Counts in expert-major iteration order.
117 counts_em = counts_2d.T.contiguous().view(-1) # [E*R]
119 # ``repeat_interleave`` expands each block's src start to one entry per
120 # token in that block — gives the source position of each output token.
121 block_src_starts = src_offsets_em.repeat_interleave(counts_em) # [total]
123 # Destination block starts in expert-major order, then expanded. The
124 # ``arange(total) - dst_block_starts_per_token`` produces 0..n-1 within
125 # each block, i.e. the intra-block offset.
126 dst_block_starts = counts_em.cumsum(0) - counts_em # [E*R]
127 dst_block_starts_per_token = dst_block_starts.repeat_interleave(counts_em)
128 intra = platform.arange(0, total, device=counts.device) - dst_block_starts_per_token
130 permuted_indices = (block_src_starts + intra).long()
131 return permuted_indices, num_tokens_per_expert
134def _permute(x, tokens_per_expert_group, ep_degree: int, num_local_experts: int):
135 """Apply rank-major → expert-major permutation to routed tokens.
137 Args:
138 x: Received token tensor of shape
139 ``[sum(tokens_per_expert_group), *feature_dims]``.
140 tokens_per_expert_group: 1-D integer tensor of shape
141 ``[ep_degree * num_local_experts]`` (output of the first
142 all-to-all that exchanges token counts).
143 ep_degree: EP group size (number of ranks).
144 num_local_experts: Number of experts owned by this rank.
146 Returns:
147 Tuple of:
149 - ``original_shape``: shape of *x* before permutation.
150 - ``permuted_x``: tokens reordered to expert-major layout.
151 - ``permuted_indices``: permutation indices (needed for
152 :func:`_unpermute`).
153 - ``num_tokens_per_expert``: token count per local expert.
154 """
155 original_shape = x.shape
156 permuted_indices, num_tokens_per_expert = _generate_permute_indices(
157 tokens_per_expert_group, num_local_experts, ep_degree
158 )
159 # ``x[permuted_indices]`` works for empty indices too (returns a
160 # shape-0 tensor with a real grad_fn). Avoid the early-return with
161 # ``new_zeros`` which would produce a leaf tensor without grad_fn and
162 # silently break autograd for ranks that happen to receive zero tokens.
163 permuted_x = x[permuted_indices]
164 return original_shape, permuted_x, permuted_indices, num_tokens_per_expert
167def _unpermute(out, original_shape, permuted_indices):
168 """Reverse the permutation applied by :func:`_permute`.
170 Args:
171 out: Expert-major output tensor of shape
172 ``[sum(num_tokens_per_expert), *feature_dims]``.
173 original_shape: Shape before permutation (from :func:`_permute`).
174 permuted_indices: Permutation indices from :func:`_permute`.
176 Returns:
177 Token tensor restored to the rank-major layout received after
178 all-to-all, with shape ``original_shape``.
179 """
180 # ``result[permuted_indices] = out`` is a differentiable scatter that
181 # also handles the empty-index case (no-op assignment, but autograd
182 # still connects ``result`` back to ``out``). Do NOT short-circuit
183 # with a bare ``new_zeros`` — that returns a leaf tensor without
184 # grad_fn and the downstream combine a2a loses its backward path,
185 # which manifests as "element 0 of tensors does not require grad".
186 result = out.new_zeros(*original_shape)
187 result[permuted_indices] = out
188 return result
191# ---------------------------------------------------------------------------
192# DispatchContext — state shared between token dispatch and combine
193# ---------------------------------------------------------------------------
195@dataclass
196class DispatchContext:
197 """Intermediate state between token dispatch and combine.
199 Stored in ``module._ep_dispatch_ctx`` for a single forward pass.
200 This solves the instance sharing problem when the same ExpertParallel
201 style object is applied to multiple layers:
203 Example problem (before this fix):
204 ep_style = ExpertParallel()
205 ep_style.apply(layer1.experts, mesh) # registers hooks
206 ep_style.apply(layer2.experts, mesh) # reuses same ep_style
208 # During forward:
209 # layer1.dispatch writes to ep_style._state_stack
210 # layer2.dispatch pushes to same stack ← INTERLEAVING
211 # layer1.combine pops wrong state (LIFO violation)
213 Solution: Store context per-module, not per-style-instance.
215 Built by :meth:`AllToAllTokenDispatcher.dispatch` and consumed by
216 :meth:`AllToAllTokenDispatcher.combine`. The caller
217 (e.g. :class:`ExpertParallel`) stores this on the module between the
218 paired dispatch/combine calls.
219 """
221 input_splits: List[int]
222 output_splits: List[int]
223 input_shape: Tuple[int, ...]
224 permuted_indices: Any
225 probs_input_shape: Optional[Tuple[int, ...]] = None
228@dataclass
229class DeredundencyDispatchContext(DispatchContext):
230 """State shared by deredundency dispatch and combine.
232 The inherited split and permutation fields describe the inner-EP
233 all-to-all. The extra fields describe the OEP shared token view and the
234 whiteboard scatter used before the final reduce-scatter combine.
235 """
237 dispatch_indices: Optional[object] = None
238 router_coeff: Optional[object] = None
239 gathered_shape: Optional[tuple] = None
240 oep_size: int = 1
244@dataclass(frozen=True)
245class _DeredundencyMeshInfo:
246 """Resolved mesh metadata for two-stage deredundency token exchange."""
248 oep_group: object
249 iep_group: object
250 oep_size: int
251 iep_size: int
252 outer_rank: int
253 inner_rank: int
256def _get_deredundency_mesh_info(device_mesh: DeviceMesh) -> _DeredundencyMeshInfo:
257 """Resolve ``oep`` / ``iep`` groups from a 1-D or 2-D EP mesh."""
258 ndim = getattr(device_mesh, "ndim", 1)
259 if not isinstance(ndim, int):
260 ndim = 1
261 if ndim == 1:
262 return _DeredundencyMeshInfo(
263 oep_group=None,
264 iep_group=device_mesh.get_group(),
265 oep_size=1,
266 iep_size=device_mesh.size(),
267 outer_rank=0,
268 inner_rank=device_mesh.get_local_rank(),
269 )
270 if ndim != 2:
271 raise ValueError(
272 "DeredundencyTokenDispatcher expects a 1-D EP mesh or a 2-D "
273 f"[oep, iep] EP mesh, but got ndim={ndim}."
274 )
276 mesh_dim_names = getattr(device_mesh, "mesh_dim_names", None) or ()
277 oep_dim = mesh_dim_names.index("oep") if "oep" in mesh_dim_names else 0
278 iep_dim = mesh_dim_names.index("iep") if "iep" in mesh_dim_names else 1
279 if oep_dim == iep_dim:
280 raise ValueError("DeredundencyTokenDispatcher requires distinct oep and iep mesh dimensions.")
282 return _DeredundencyMeshInfo(
283 oep_group=device_mesh.get_group(oep_dim),
284 iep_group=device_mesh.get_group(iep_dim),
285 oep_size=device_mesh.size(oep_dim),
286 iep_size=device_mesh.size(iep_dim),
287 outer_rank=device_mesh.get_local_rank(oep_dim),
288 inner_rank=device_mesh.get_local_rank(iep_dim),
289 )
292def _generate_deredundency_dispatch_indices(
293 tokens_per_expert_by_source,
294 expert_start: int,
295 iep_size: int,
296 num_local_experts: int,
297):
298 """Generate gather-view indices ordered by IEP destination rank.
300 ``tokens_per_expert_by_source`` is shaped ``[oep_size, num_experts]`` and
301 describes each source rank's expert-major routed buffer after the OEP
302 all-gather. The returned indices select the current outer expert range
303 and order it as ``[iep_dst, local_expert, oep_source]`` so each IEP
304 destination chunk keeps local-expert blocks contiguous for the later
305 rank-major → expert-major permutation.
306 """
307 oep_size = tokens_per_expert_by_source.shape[0]
308 experts_per_outer = iep_size * num_local_experts
309 expert_end = expert_start + experts_per_outer
311 source_totals = tokens_per_expert_by_source.sum(dim=1)
312 source_offsets = source_totals.cumsum(0) - source_totals
313 expert_offsets = (
314 tokens_per_expert_by_source.cumsum(dim=1)
315 - tokens_per_expert_by_source
316 + source_offsets.view(oep_size, 1)
317 )
319 selected_counts = tokens_per_expert_by_source[:, expert_start:expert_end].view(
320 oep_size, iep_size, num_local_experts,
321 )
322 selected_offsets = expert_offsets[:, expert_start:expert_end].view(
323 oep_size, iep_size, num_local_experts,
324 )
325 counts_by_destination = selected_counts.permute(1, 2, 0).contiguous()
326 offsets_by_destination = selected_offsets.permute(1, 2, 0).contiguous()
328 block_counts = counts_by_destination.view(-1)
329 token_counts_by_destination_expert = selected_counts.sum(dim=0).contiguous().view(-1)
330 total = int(block_counts.sum())
331 if total == 0:
332 return block_counts.new_zeros(0, dtype=block_counts.dtype).long(), token_counts_by_destination_expert
334 block_starts = offsets_by_destination.view(-1).repeat_interleave(block_counts)
335 block_offsets = block_counts.cumsum(0) - block_counts
336 block_offsets_per_token = block_offsets.repeat_interleave(block_counts)
337 intra = platform.arange(0, total, device=tokens_per_expert_by_source.device) - block_offsets_per_token
339 return (block_starts + intra).long(), token_counts_by_destination_expert
342def _scale_by_router_coeff(tokens, router_coeff):
343 """Scale routed expert outputs by optional router coefficients."""
344 if router_coeff is None:
345 return tokens
346 if router_coeff.shape[0] != tokens.shape[0]:
347 raise ValueError(
348 "router_coeff length must match routed token count, got "
349 f"{router_coeff.shape[0]} and {tokens.shape[0]}."
350 )
351 coeff = router_coeff
352 if len(coeff.shape) == 1 and len(tokens.shape) > 1:
353 coeff = coeff.reshape((-1,) + (1,) * (len(tokens.shape) - 1))
354 return tokens * coeff
357def _scatter_add_first_dim(src, indices, output_shape):
358 """Scatter-add rows of ``src`` into a zero tensor along dim 0."""
359 result = src.new_zeros(*output_shape)
360 if len(src.shape) == 1:
361 scatter_indices = indices
362 else:
363 scatter_indices = indices.reshape((-1,) + (1,) * (len(src.shape) - 1)).expand(
364 -1, *src.shape[1:],
365 )
366 if hasattr(result, "scatter_add"):
367 return result.scatter_add(0, scatter_indices, src)
368 if hasattr(result, "index_add"):
369 return result.index_add(0, indices, src)
370 raise RuntimeError(
371 "DeredundencyTokenDispatcher.combine requires tensor scatter_add or "
372 "index_add support for exdispatch_idx accumulation."
373 )
376class _DeredundencyCombineHandle(AsyncHandle):
377 """Async handle that finishes deredundency combine post-processing."""
379 def __init__(
380 self,
381 async_tensor: object,
382 mesh_info: _DeredundencyMeshInfo,
383 ctx: DeredundencyDispatchContext,
384 ) -> None:
385 super().__init__(async_tensor)
386 self._mesh_info = mesh_info
387 self._ctx = ctx
388 self._combined: Optional[object] = None
390 def wait(self) -> object:
391 """Wait for IEP a2a, then finish OEP scatter/reduce combine once."""
392 if self._combined is None:
393 outer_output = super().wait()
394 weighted_output = _scale_by_router_coeff(outer_output, self._ctx.router_coeff)
395 combine_whiteboard = _scatter_add_first_dim(
396 weighted_output,
397 self._ctx.dispatch_indices,
398 self._ctx.gathered_shape,
399 )
400 if self._ctx.oep_size == 1:
401 self._combined = combine_whiteboard
402 else:
403 self._combined = platform.differentiable_reduce_scatter(
404 combine_whiteboard,
405 self._ctx.oep_size,
406 0,
407 "sum",
408 self._mesh_info.oep_group,
409 )
410 return self._combined
413# ---------------------------------------------------------------------------
414# BaseExpertParallel — abstract base for all-to-all EP strategies
415# ---------------------------------------------------------------------------
417class BaseExpertParallel(ParallelStyle, ABC):
418 """Abstract base class for Expert Parallel strategies with token dispatch.
420 Subclasses implement :meth:`_partition_fn`, :meth:`_token_dispatch`, and
421 :meth:`_token_combine`; this class wires them into :func:`distribute_module`.
422 """
424 def apply(self, module: Module, device_mesh: DeviceMesh) -> Module:
425 """Apply EP sharding and dispatch/combine hooks to *module*.
427 Args:
428 module: A :class:`~hyper_parallel.core.expert_parallel.moe.GroupedExperts`
429 instance to shard.
430 device_mesh: Device mesh for this EP strategy.
432 Returns:
433 The module with distributed parameters and dispatch/combine hooks.
434 """
435 return distribute_module(
436 module,
437 device_mesh,
438 self._partition_fn,
439 self._token_dispatch,
440 self._token_combine,
441 )
443 @abstractmethod
444 def _partition_fn(
445 self, name: str, module: Module, device_mesh: DeviceMesh
446 ) -> None:
447 """Shard module parameters according to this strategy.
449 Args:
450 name: Submodule name.
451 module: The module whose parameters are being sharded.
452 device_mesh: Device mesh for this EP strategy.
453 """
455 @abstractmethod
456 def _token_dispatch(self, module: Module, inputs, device_mesh: DeviceMesh):
457 """Pre-hook: route input tokens to their assigned ranks.
459 Args:
460 module: The ``GroupedExperts`` module.
461 inputs: Forward inputs tuple.
462 device_mesh: Device mesh for this EP strategy.
464 Returns:
465 Transformed inputs for local expert computation.
466 """
468 @abstractmethod
469 def _token_combine(self, module: Module, routed_output, device_mesh: DeviceMesh):
470 """Post-hook: gather expert outputs back to the originating ranks.
472 Args:
473 module: The ``GroupedExperts`` module.
474 routed_output: Expert output tensor in expert-major order.
475 device_mesh: Device mesh for this EP strategy.
477 Returns:
478 Token tensor in the original token-major layout.
479 """
482# ---------------------------------------------------------------------------
483# AllToAllTokenDispatcher — token dispatch/combine via all-to-all
484# ---------------------------------------------------------------------------
486class AllToAllTokenDispatcher:
487 """Token dispatch and combine via all-to-all for expert parallelism.
489 Provides :meth:`dispatch` and :meth:`combine` as static methods that
490 receive and return a :class:`DispatchContext` object. This decouples
491 the all-to-all token routing logic from the parallel style class so
492 that it can be reused or tested independently.
494 Callers (e.g. :class:`ExpertParallel`) are responsible for storing the
495 context between the paired dispatch/combine calls.
496 """
498 @staticmethod
499 def dispatch(module: Module, inputs: tuple, device_mesh: DeviceMesh) -> tuple:
500 """Dispatch tokens to their assigned ranks via all-to-all.
502 Called as an ``input_fn`` hook by :func:`distribute_module`. Receives
503 the module's forward inputs and returns transformed inputs.
505 Args:
506 module: The ``GroupedExperts`` module.
507 inputs: Tuple ``(routed_input, num_tokens_per_expert)`` or
508 ``(routed_input, num_tokens_per_expert, permuted_probs)`` where
509 ``routed_input`` has shape ``[total_tokens, dim]``, and
510 ``num_tokens_per_expert`` has shape ``[num_experts]``, and
511 ``permuted_probs`` has shape ``[total_tokens]``.
512 device_mesh: EP device mesh (1-D).
514 Returns:
515 Tuple ``(permuted_local_input, local_token_counts, ctx)`` or
516 ``(permuted_local_input, local_token_counts, permuted_probs, ctx)``
517 depending on whether *permuted_probs* was provided -
518 the first two elements are the transformed inputs for local
519 expert computation; *ctx* is a :class:`DispatchContext`
520 carrying the updated state to be stored by the caller.
521 """
522 del module # module unused, kept for API consistency
523 routed_input, num_tokens_per_expert = inputs[0], inputs[1]
524 permuted_probs = inputs[2] if len(inputs) > 2 else None
525 ep_group = device_mesh.get_group()
526 ep_size = device_mesh.size()
527 num_local_experts = num_tokens_per_expert.shape[0] // ep_size
529 # --- Step 1: exchange token counts (no gradient needed) ---
530 # Each rank needs to know how many tokens it will receive from every
531 # other rank (for each local expert). Uses ``async_op=True`` + an
532 # explicit ``handle.wait()`` rather than ``async_op=False`` because
533 # the implicit cross-stream sync is NCCL-only; on HCCL the compute
534 # stream may read ``counts_out`` before the collective write is
535 # visible, producing garbage values that blow up the downstream
536 # ``torch.empty(sum(output_splits), ...)`` allocation.
537 counts_out, handle = platform.all_to_all_single(
538 num_tokens_per_expert,
539 output_shape=[num_tokens_per_expert.shape[0]],
540 group=ep_group,
541 async_op=True,
542 )
543 if handle is not None:
544 handle.wait()
545 # counts_out shape: [ep_size * num_local_experts]
546 # counts_out[r * num_local_experts + e] = tokens from rank r for expert e
548 # --- Step 2: compute input / output splits ---
549 # input_splits[r] = tokens this rank sends to rank r
550 # output_splits[r] = tokens this rank receives from rank r
551 # Reshape to [ep_size, num_local_experts] and sum per rank on device;
552 # a single ``tolist()`` drains the rank-sum vector to host, replacing
553 # ``2 * ep_size`` scalar ``int()`` D2H syncs with 2.
554 input_splits = num_tokens_per_expert.view(ep_size, num_local_experts).sum(dim=1).tolist()
555 output_splits = counts_out.view(ep_size, num_local_experts).sum(dim=1).tolist()
557 # --- Step 3a: exchange actual tokens (differentiable) ---
558 dispatched = platform.differentiable_all_to_all_single(
559 routed_input, input_splits, output_splits, group=ep_group,
560 )
562 # --- Step 4a: rank-major → expert-major permutation ---
563 input_shape, permuted, permuted_indices, local_counts = _permute(
564 dispatched, counts_out, ep_size, num_local_experts
565 )
567 # Steps 3b + 4b: exchange and reorder permuted_probs (when provided)
568 # alongside the tokens, keeping weights aligned with the expert-major
569 # token order. Extracted so dispatch stays under the Lizard NLOC cap.
570 probs_input_shape, permuted_probs_reordered = (
571 AllToAllTokenDispatcher._exchange_and_reorder_probs(
572 permuted_probs, input_splits, output_splits,
573 counts_out, ep_size, num_local_experts, ep_group,
574 )
575 )
577 # Build dispatch context for combine step.
578 # Caller (e.g., ExpertParallel._token_dispatch) is responsible for storing
579 # this context and passing it to combine(). This decouples dispatch/combine
580 # from module state and solves the instance sharing problem.
581 ctx = DispatchContext(
582 input_splits=input_splits,
583 output_splits=output_splits,
584 input_shape=input_shape,
585 permuted_indices=permuted_indices,
586 probs_input_shape=probs_input_shape,
587 )
589 if permuted_probs is not None:
590 return permuted, local_counts, permuted_probs_reordered, ctx
591 return permuted, local_counts, ctx
593 @staticmethod
594 def _exchange_and_reorder_probs(permuted_probs, input_splits, output_splits,
595 counts_out, ep_size, num_local_experts, ep_group):
596 """Exchange ``permuted_probs`` via a2a and reorder to expert-major.
598 Returns ``(probs_input_shape, permuted_probs_reordered)``; both are
599 ``None`` when *permuted_probs* is ``None`` (no-op, no second a2a).
600 """
601 if permuted_probs is None:
602 return None, None
603 # --- Step 3b: exchange permuted_probs via all-to-all (differentiable) ---
604 dispatched_probs = platform.differentiable_all_to_all_single(
605 permuted_probs, input_splits, output_splits, group=ep_group,
606 )
607 # --- Step 4b: rank-major → expert-major permutation for probs ---
608 probs_input_shape, permuted_probs_reordered, _, _ = _permute(
609 dispatched_probs, counts_out, ep_size, num_local_experts
610 )
611 return probs_input_shape, permuted_probs_reordered
613 @staticmethod
614 def combine(module: Module, routed_output: object, device_mesh: DeviceMesh, ctx: DispatchContext) -> object:
615 """Gather expert outputs back to the originating ranks via all-to-all.
617 Called as an ``output_fn`` hook by :func:`distribute_module`.
618 Receives dispatch context from the caller (previously returned by dispatch).
620 Args:
621 module: The ``GroupedExperts`` module (unused, for API consistency).
622 routed_output: Expert output tensor in expert-major order,
623 shape ``[sum(local_counts), dim]``.
624 device_mesh: EP device mesh (1-D).
625 ctx: :class:`DispatchContext` previously returned by
626 :meth:`dispatch`.
628 Returns:
629 Token tensor in the original token-major layout,
630 shape ``[sum(input_splits), dim]``.
631 """
632 del module # module not used, kept for API consistency
633 ep_group = device_mesh.get_group()
635 # expert-major → rank-major
636 unpermuted = _unpermute(routed_output, ctx.input_shape, ctx.permuted_indices)
638 # reverse all-to-all (output/input splits are swapped)
639 combined = platform.differentiable_all_to_all_single(
640 unpermuted,
641 ctx.output_splits, # was output, now becomes input
642 ctx.input_splits, # was input, now becomes output
643 group=ep_group,
644 )
645 return combined
647 @staticmethod
648 def combine_start(routed_output, device_mesh, ctx):
649 """Launch async combine all-to-all without waiting for completion.
651 Splits the combine into two phases so that the caller can overlap
652 the a2a communication with independent computation (e.g. a shared
653 expert forward pass). The caller must later call
654 :meth:`combine_wait` or ``handle.wait()`` to obtain the final
655 result.
657 Step 1 (synchronous, local): expert-major → rank-major unpermute.
658 Step 2 (asynchronous, cross-rank): reverse all-to-all.
660 Args:
661 routed_output: Expert output tensor in expert-major order,
662 shape ``[sum(local_counts), dim]``.
663 device_mesh: EP device mesh (1-D).
664 ctx: :class:`DispatchContext` previously returned by
665 :meth:`dispatch`.
667 Returns:
668 :class:`AsyncHandle` carrying the state needed by
669 :meth:`combine_wait`.
670 """
671 ep_group = device_mesh.get_group()
673 # expert-major → rank-major (local, no communication)
674 unpermuted = _unpermute(routed_output, ctx.input_shape, ctx.permuted_indices)
676 # async reverse all-to-all (output/input splits are swapped)
677 combined_async = platform.differentiable_all_to_all_single_async(
678 unpermuted,
679 ctx.output_splits,
680 ctx.input_splits,
681 group=ep_group,
682 )
684 return AsyncHandle(combined_async)
686 @staticmethod
687 def combine_wait(handle):
688 """Wait for the async combine all-to-all to complete.
690 Args:
691 handle: :class:`AsyncHandle` returned by :meth:`combine_start`.
693 Returns:
694 Combined tensor in the original token-major layout.
695 """
696 return handle.wait()
699# ---------------------------------------------------------------------------
700# DeredundencyTokenDispatcher — token dispatch via OEP all-gather + IEP all-to-all
701# ---------------------------------------------------------------------------
703class DeredundencyTokenDispatcher:
704 """Token dispatch/combine via OEP all-gather plus IEP all-to-all.
706 This dispatcher keeps the same public contract as
707 :class:`AllToAllTokenDispatcher`, but decomposes the global EP all-to-all
708 into the deredundency flow described in
709 ``docs/moe_alltoall_deredundency_token_permutation.md``:
711 1. Form a shared token/count view across the OEP group.
712 2. Select only the current outer expert range.
713 3. Send selected tokens to concrete local-expert ranks inside the IEP
714 group.
715 4. Sort received tokens into local expert-major order.
717 For a 2-D mesh, dimension ``"oep"`` / ``0`` is the outer group and
718 ``"iep"`` / ``1`` is the inner group. A 1-D mesh is treated as
719 ``oep_size == 1`` and degenerates to the standard all-to-all data flow.
720 """
722 @staticmethod
723 def _oep_gather_for_dispatch(
724 num_tokens_per_expert,
725 routed_input,
726 router_coeff,
727 mesh_info: _DeredundencyMeshInfo,
728 ) -> tuple:
729 """All-gather token counts and routed input across the OEP group.
731 Args:
732 num_tokens_per_expert: Token count per expert ``[num_experts]``.
733 routed_input: Routed token tensor ``[total_tokens, dim]``.
734 router_coeff: Optional router coefficients ``[total_tokens]``.
735 mesh_info: Resolved OEP/IEP mesh descriptor.
737 Returns:
738 Tuple ``(gathered_counts, gathered_routed, gathered_router_coeff)``
739 where ``gathered_counts`` has shape ``[oep_size, num_experts]``,
740 ``gathered_routed`` has shape ``[oep_size * total_tokens, dim]``,
741 and ``gathered_router_coeff`` is the gathered coefficients or None.
743 Raises:
744 ValueError: If routed token counts differ across OEP ranks.
745 """
746 if mesh_info.oep_size == 1:
747 gathered_counts = num_tokens_per_expert.view(1, num_tokens_per_expert.shape[0])
748 return gathered_counts, routed_input, router_coeff
750 gathered_counts, handle = platform.all_gather_single(
751 num_tokens_per_expert,
752 output_shape=[mesh_info.oep_size * num_tokens_per_expert.shape[0]],
753 group=mesh_info.oep_group,
754 async_op=True,
755 )
756 if handle is not None:
757 handle.wait()
758 gathered_counts = gathered_counts.view(mesh_info.oep_size, num_tokens_per_expert.shape[0])
759 source_token_totals = gathered_counts.sum(dim=1).tolist()
760 if any(total != routed_input.shape[0] for total in source_token_totals):
761 raise ValueError(
762 "DeredundencyTokenDispatcher requires equal routed token "
763 "counts within each OEP group because the shared token view "
764 f"uses all-gather, got totals {source_token_totals}."
765 )
766 gathered_routed = platform.differentiable_all_gather_concat(
767 routed_input, mesh_info.oep_group, mesh_info.oep_size, 0,
768 )
769 if router_coeff is None:
770 gathered_router_coeff = None
771 else:
772 gathered_router_coeff = platform.differentiable_all_gather_concat(
773 router_coeff, mesh_info.oep_group, mesh_info.oep_size, 0,
774 )
775 return gathered_counts, gathered_routed, gathered_router_coeff
777 @staticmethod
778 def dispatch(module: Module, inputs: tuple, device_mesh: DeviceMesh) -> tuple:
779 """Dispatch tokens using OEP all-gather and IEP all-to-all.
781 Args:
782 module: The ``GroupedExperts`` module (unused here).
783 inputs: Tuple ``(routed_input, num_tokens_per_expert)`` or
784 ``(routed_input, num_tokens_per_expert, router_coeff)`` where
785 ``routed_input`` has shape ``[total_tokens, dim]``,
786 ``num_tokens_per_expert`` has shape ``[num_experts]``, and
787 ``router_coeff`` is an optional 1-D tensor of shape
788 ``[total_tokens]``.
789 device_mesh: 1-D EP mesh or 2-D ``[oep, iep]`` EP mesh.
791 Returns:
792 Tuple ``(permuted_local_input, local_token_counts, ctx)`` with the
793 same meaning as :meth:`AllToAllTokenDispatcher.dispatch`. Unlike
794 :meth:`AllToAllTokenDispatcher.dispatch`, this always returns a
795 3-tuple — ``router_coeff`` (when provided) is consumed by the
796 combine step, not returned for in-expert weighting.
798 Raises:
799 ValueError: If the expert count is not divisible by the full EP
800 size represented by the deredundency mesh.
802 Note:
803 The third input is **router_coeff**, *not* ``permuted_probs``.
804 Both occupy ``inputs[2]`` with shape ``[total_tokens]`` but differ:
805 ``permuted_probs`` (:class:`AllToAllTokenDispatcher`) is exchanged
806 with tokens and applied **inside** experts (pre-w2 activation,
807 i.e. ``score_before_experts=False``); ``router_coeff`` (here) is
808 OEP-gathered, scattered with tokens, and applied in
809 :meth:`combine` to the expert **final output**. The w2/SiLU
810 nonlinearity makes the two mathematically distinct, so
811 deredundency does **not** support ``score_before_experts=False``
812 — passing scores as ``inputs[2]`` would be silently treated as
813 ``router_coeff``. Use ``score_before_experts=True`` until
814 deredundency's weighting is redesigned.
815 """
816 del module
817 routed_input, num_tokens_per_expert = inputs[0], inputs[1]
818 router_coeff = inputs[2] if len(inputs) > 2 else None
819 if router_coeff is not None and router_coeff.shape[0] != routed_input.shape[0]:
820 raise ValueError(
821 "router_coeff length must match routed_input token count, got "
822 f"{router_coeff.shape[0]} and {routed_input.shape[0]}."
823 )
824 mesh_info = _get_deredundency_mesh_info(device_mesh)
825 ep_size = mesh_info.oep_size * mesh_info.iep_size
826 if num_tokens_per_expert.shape[0] % ep_size != 0:
827 raise ValueError(
828 "num_tokens_per_expert length must be divisible by the full "
829 f"EP size {ep_size}, got {num_tokens_per_expert.shape[0]}."
830 )
831 num_local_experts = num_tokens_per_expert.shape[0] // ep_size
832 experts_per_outer = mesh_info.iep_size * num_local_experts
833 expert_start = mesh_info.outer_rank * experts_per_outer
835 gathered_counts, gathered_routed, gathered_router_coeff = (
836 DeredundencyTokenDispatcher._oep_gather_for_dispatch(
837 num_tokens_per_expert, routed_input, router_coeff, mesh_info,
838 )
839 )
841 dispatch_indices, node_counts_per_expert = _generate_deredundency_dispatch_indices(
842 gathered_counts,
843 expert_start,
844 mesh_info.iep_size,
845 num_local_experts,
846 )
847 iep_input_splits = node_counts_per_expert.view(
848 mesh_info.iep_size, num_local_experts).sum(dim=1).tolist()
849 iep_counts_out, iep_output_splits = (
850 DeredundencyTokenDispatcher._iep_exchange_counts(
851 node_counts_per_expert, mesh_info, num_local_experts,
852 )
853 )
854 outer_routed_input = gathered_routed[dispatch_indices]
855 outer_router_coeff = (
856 None if gathered_router_coeff is None else gathered_router_coeff[dispatch_indices]
857 )
858 dispatched = platform.differentiable_all_to_all_single(
859 outer_routed_input, iep_input_splits, iep_output_splits,
860 group=mesh_info.iep_group,
861 )
862 input_shape, permuted, permuted_indices, local_counts = _permute(
863 dispatched, iep_counts_out, mesh_info.iep_size, num_local_experts,
864 )
865 ctx = DeredundencyDispatchContext(
866 input_splits=iep_input_splits,
867 output_splits=iep_output_splits,
868 input_shape=input_shape,
869 permuted_indices=permuted_indices,
870 dispatch_indices=dispatch_indices,
871 router_coeff=outer_router_coeff,
872 gathered_shape=gathered_routed.shape,
873 oep_size=mesh_info.oep_size,
874 )
875 return permuted, local_counts, ctx
877 @staticmethod
878 def _iep_exchange_counts(node_counts_per_expert, mesh_info, num_local_experts):
879 """IEP all-to-all of per-expert counts; return (counts_out, output_splits)."""
880 iep_counts_out, handle = platform.all_to_all_single(
881 node_counts_per_expert,
882 output_shape=[node_counts_per_expert.shape[0]],
883 group=mesh_info.iep_group, async_op=True,
884 )
885 if handle is not None:
886 handle.wait()
887 iep_output_splits = iep_counts_out.view(
888 mesh_info.iep_size, num_local_experts).sum(dim=1).tolist()
889 return iep_counts_out, iep_output_splits
891 @staticmethod
892 def combine(module: Module, routed_output: object, device_mesh: DeviceMesh,
893 ctx: DeredundencyDispatchContext) -> object:
894 """Gather expert outputs back to the originating ranks.
896 Args:
897 module: The ``GroupedExperts`` module (unused).
898 routed_output: Expert output tensor in expert-major order.
899 device_mesh: 1-D EP mesh or 2-D ``[oep, iep]`` EP mesh.
900 ctx: Context returned by :meth:`dispatch`.
902 Returns:
903 Token tensor in the original source-rank routed order.
904 """
905 del module
906 mesh_info = _get_deredundency_mesh_info(device_mesh)
907 DeredundencyTokenDispatcher._validate_combine_mesh(mesh_info, ctx)
909 unpermuted = _unpermute(routed_output, ctx.input_shape, ctx.permuted_indices)
910 outer_output = platform.differentiable_all_to_all_single(
911 unpermuted,
912 ctx.output_splits,
913 ctx.input_splits,
914 group=mesh_info.iep_group,
915 )
917 weighted_output = _scale_by_router_coeff(outer_output, ctx.router_coeff)
918 combine_whiteboard = _scatter_add_first_dim(
919 weighted_output, ctx.dispatch_indices, ctx.gathered_shape,
920 )
921 if ctx.oep_size == 1:
922 return combine_whiteboard
924 return platform.differentiable_reduce_scatter(
925 combine_whiteboard,
926 ctx.oep_size,
927 0,
928 "sum",
929 mesh_info.oep_group,
930 )
932 @staticmethod
933 def _validate_combine_mesh(
934 mesh_info: _DeredundencyMeshInfo,
935 ctx: DeredundencyDispatchContext,
936 ) -> None:
937 """Validate that dispatch context and combine mesh are compatible."""
938 if mesh_info.oep_size != ctx.oep_size:
939 raise ValueError(
940 "DeredundencyTokenDispatcher.combine received a context for "
941 f"oep_size={ctx.oep_size}, but the mesh resolves to oep_size={mesh_info.oep_size}."
942 )
944 @staticmethod
945 def combine_start(
946 routed_output: object,
947 device_mesh: DeviceMesh,
948 ctx: DeredundencyDispatchContext,
949 ) -> AsyncHandle:
950 """Launch async IEP combine all-to-all and defer deredundency post-processing.
952 The local expert-major → rank-major unpermute is performed
953 synchronously. The reverse IEP all-to-all is launched asynchronously,
954 and :meth:`combine_wait` finishes router weighting, whiteboard
955 scatter-add, and optional OEP reduce-scatter after the async output is
956 materialised.
958 Args:
959 routed_output: Expert output tensor in expert-major order.
960 device_mesh: 1-D EP mesh or 2-D ``[oep, iep]`` EP mesh.
961 ctx: Context returned by :meth:`dispatch`.
963 Returns:
964 :class:`AsyncHandle` carrying the pending IEP a2a and deredundency
965 combine state.
966 """
967 mesh_info = _get_deredundency_mesh_info(device_mesh)
968 DeredundencyTokenDispatcher._validate_combine_mesh(mesh_info, ctx)
970 unpermuted = _unpermute(routed_output, ctx.input_shape, ctx.permuted_indices)
971 outer_output_async = platform.differentiable_all_to_all_single_async(
972 unpermuted,
973 ctx.output_splits,
974 ctx.input_splits,
975 group=mesh_info.iep_group,
976 )
977 return _DeredundencyCombineHandle(outer_output_async, mesh_info, ctx)
979 @staticmethod
980 def combine_wait(handle: AsyncHandle) -> object:
981 """Wait for async deredundency combine and return the final tensor.
983 Args:
984 handle: :class:`AsyncHandle` returned by :meth:`combine_start`.
986 Returns:
987 Token tensor in the original source-rank routed order.
988 """
989 return handle.wait()
992_TOKEN_DISPATCHERS = {
993 "all_to_all": AllToAllTokenDispatcher,
994 "deredundency": DeredundencyTokenDispatcher,
995}
998def _resolve_token_dispatcher(token_dispatcher: str):
999 """Resolve a token dispatcher name to its implementation class."""
1000 try:
1001 return _TOKEN_DISPATCHERS[token_dispatcher]
1002 except KeyError as exc:
1003 supported = "', '".join(sorted(_TOKEN_DISPATCHERS))
1004 raise ValueError(
1005 f"token_dispatcher must be one of '{supported}', got {token_dispatcher!r}."
1006 ) from exc
1009def _get_flattened_ep_mesh(device_mesh: DeviceMesh) -> DeviceMesh:
1010 """Return a 1-D EP mesh, flattening a 2-D deredundency mesh if needed."""
1011 if getattr(device_mesh, "ndim", 1) == 1:
1012 return device_mesh
1013 mesh_dim_names = getattr(device_mesh, "mesh_dim_names", None) or ()
1014 if "ep" in mesh_dim_names or "ep" in device_mesh.get_flatten_mapping():
1015 return device_mesh["ep"]
1016 if set(mesh_dim_names) == {"oep", "iep"}:
1017 return device_mesh.flatten("ep")
1018 raise ValueError(
1019 "Deredundency ExpertParallel expects a 1-D EP mesh or a 2-D "
1020 "[oep, iep] mesh when partitioning expert weights."
1021 )
1024# ---------------------------------------------------------------------------
1025# ExpertParallel — standard all-to-all EP
1026# ---------------------------------------------------------------------------
1028class ExpertParallel(BaseExpertParallel):
1029 """Expert Parallel: shard experts across ranks via all-to-all token routing.
1031 Applies :meth:`apply` to a :class:`GroupedExperts` module:
1033 1. **Partition** — distributes expert weights on dim 0 (``Shard(0)``) so
1034 each rank holds ``num_experts // ep_degree`` local experts.
1035 2. **Token dispatch** (forward pre-hook) — two-step all-to-all:
1036 a. Exchange token counts (non-differentiable).
1037 b. Exchange actual tokens (differentiable, gradient flows back).
1038 Followed by rank-major → expert-major permutation.
1039 3. **Token combine** (forward post-hook) — expert-major → rank-major
1040 unpermute, then reverse all-to-all (differentiable).
1042 All collectives use ``platform.differentiable_all_to_all_single`` /
1043 ``platform.all_to_all_single`` — no direct ``torch.distributed`` calls.
1045 The token dispatcher is selectable. ``"all_to_all"`` uses
1046 :class:`AllToAllTokenDispatcher`; ``"deredundency"`` uses
1047 :class:`DeredundencyTokenDispatcher`.
1049 Args:
1050 token_dispatcher: Token dispatch strategy. Supported values are
1051 ``"all_to_all"`` and ``"deredundency"``.
1052 async_combine: When ``True``, the combine all-to-all is launched
1053 asynchronously so that the caller (e.g. :class:`MoE`) can
1054 overlap it with shared-expert computation. When ``False``
1055 (default), combine is fully synchronous — no overlap, identical
1056 to the baseline.
1058 Example::
1059 >>> ep_style = ExpertParallel()
1060 >>> sharded_experts = ep_style.apply(experts_module, ep_device_mesh)
1061 >>> # With async combine for shared-expert overlap:
1062 >>> ep_style = ExpertParallel(async_combine=True)
1063 >>> sharded_experts = ep_style.apply(experts_module, ep_device_mesh)
1064 """
1066 def __init__(self, token_dispatcher: Union[str, bool] = "all_to_all", async_combine: bool = False) -> None:
1067 """Initialize ExpertParallel.
1069 Args:
1070 token_dispatcher: Token dispatch strategy. Supported values are
1071 ``"all_to_all"`` and ``"deredundency"``.
1072 async_combine: If ``True``, use asynchronous combine all-to-all
1073 to overlap communication with shared-expert computation.
1074 """
1075 if isinstance(token_dispatcher, bool):
1076 async_combine = token_dispatcher
1077 token_dispatcher = "all_to_all"
1078 self._dispatch_ctx: Optional[DispatchContext] = None
1079 self.async_combine = async_combine
1080 self._token_dispatcher_name = token_dispatcher
1081 self._token_dispatcher = _resolve_token_dispatcher(token_dispatcher)
1083 def _token_dispatch(self, module: Module, inputs, device_mesh: DeviceMesh):
1084 """Dispatch tokens to their assigned ranks via all-to-all.
1086 Delegates to the configured token dispatcher and stores the
1087 returned :class:`DispatchContext` on the instance for the matching
1088 :meth:`_token_combine` call.
1090 Args:
1091 module: The ``GroupedExperts`` module.
1092 inputs: Tuple ``(routed_input, num_tokens_per_expert)`` or
1093 ``(routed_input, num_tokens_per_expert, routed_probs)``.
1094 device_mesh: EP device mesh (1-D).
1096 Returns:
1097 Tuple ``(permuted_local_input, local_token_counts)`` or
1098 ``(permuted_local_input, local_token_counts, permuted_probs)``
1099 depending on whether *routed_probs* was provided.
1100 """
1101 # Delegate to the configured dispatcher (all_to_all or deredundency).
1102 # Hard-coding AllToAllTokenDispatcher here would mismatch _token_combine
1103 # (which uses self._token_dispatcher) and break deredundency.
1104 dispatch_result = self._token_dispatcher.dispatch(module, inputs, device_mesh)
1105 ctx = dispatch_result[-1]
1106 # Store context in module attribute for _token_combine to read.
1107 # Using module attribute ensures each module has its own context,
1108 # solving the instance sharing problem when the same ExpertParallel
1109 # style object is applied to multiple GroupedExperts modules.
1110 # pylint: disable=W0212
1111 module._ep_dispatch_ctx = ctx
1112 # dispatch_result is either (permuted, local_counts, ctx) or
1113 # (permuted, local_counts, permuted_probs, ctx); return all but ctx.
1114 return dispatch_result[:-1]
1116 def _token_combine(self, module: Module, routed_output, device_mesh: DeviceMesh):
1117 """Gather expert outputs back to the originating ranks via all-to-all.
1119 When ``async_combine=True``, launches the combine all-to-all
1120 asynchronously and returns an :class:`AsyncCollectiveTensor`. The
1121 actual device-side wait is deferred until the downstream consumer
1122 (e.g. MoE unpermutation) first reads the tensor, enabling overlap
1123 with shared-expert computation.
1125 When ``async_combine=False`` (default), uses the synchronous
1126 :meth:`AllToAllTokenDispatcher.combine` — identical to the baseline.
1128 Args:
1129 module: The ``GroupedExperts`` module.
1130 routed_output: Expert output tensor in expert-major order.
1131 device_mesh: EP device mesh (1-D).
1133 Returns:
1134 Token tensor in the original token-major layout. When
1135 ``async_combine=True``, this may be an async collective tensor
1136 whose values are not yet materialised.
1138 Raises:
1139 RuntimeError: If dispatch context is not found (dispatch was not called).
1140 """
1141 # Read dispatch context from module attribute set by _token_dispatch.
1142 # pylint: disable=W0212
1143 ctx = getattr(module, "_ep_dispatch_ctx", None)
1144 if ctx is None:
1145 raise RuntimeError(
1146 "_token_combine called but no dispatch context found in module. "
1147 "This indicates _token_dispatch was not called before _token_combine, "
1148 "or the context was already consumed by a previous combine call."
1149 )
1151 # Note: Do NOT delete the context here. In PyTorch, the tensors in ctx
1152 # are captured by autograd graph and don't need the attribute. But in
1153 # MindSpore PyNative mode, deleting the attribute may break backward.
1154 # The context will be overwritten on the next forward call.
1156 if self.async_combine:
1157 handle = self._token_dispatcher.combine_start(
1158 routed_output, device_mesh, ctx
1159 )
1160 # Store on module for external inspection / advanced use cases.
1161 # pylint: disable=W0212
1162 module._ep_combine_handle = handle
1163 # Return the async tensor. The first non-view access by the
1164 # downstream consumer (e.g. MoE unpermutation) will trigger the
1165 # implicit wait, overlapping with shared_expert computation.
1166 return handle.wait()
1168 return self._token_dispatcher.combine(
1169 module, routed_output, device_mesh, ctx,
1170 )
1172 def _partition_mesh(self, device_mesh: DeviceMesh) -> DeviceMesh:
1173 """Return the mesh used to shard expert weights."""
1174 if self._token_dispatcher_name == "deredundency":
1175 return _get_flattened_ep_mesh(device_mesh)
1176 return device_mesh
1178 def _partition_fn(
1179 self, name: str, module: Module, device_mesh: DeviceMesh
1180 ) -> None:
1181 """Shard all expert parameters along dim 0 (expert dimension).
1183 Args:
1184 name: Submodule name (unused).
1185 module: The module whose parameters are being sharded.
1186 device_mesh: EP device mesh.
1187 """
1188 del name
1189 partition_mesh = self._partition_mesh(device_mesh)
1190 for key, param in _distribute_module_iter_params(module):
1191 if param is None:
1192 continue
1193 src = _distribute_module_param_source(param)
1194 requires_grad = bool(getattr(param, "requires_grad", True))
1195 dt = distribute_tensor(src, partition_mesh, [Shard(0)])
1196 new_param = _distribute_module_new_parameter(key, dt, requires_grad)
1197 _distribute_module_set_param(module, key, new_param)
1200# ---------------------------------------------------------------------------
1201# TensorParallel — TP-only weight sharding for experts (no token dispatch)
1202# ---------------------------------------------------------------------------
1203class TensorParallel(ParallelStyle):
1204 """Tensor Parallel for expert weights (no token dispatch).
1206 Shards the ``GroupedExperts`` weight tensors in the column/row-wise
1207 pattern used by standard TP:
1209 - ``w1`` / ``w3``: ``Shard(1)`` — column-wise (hidden_dim dimension).
1210 - ``w2``: ``Shard(2)`` — row-wise (output dim dimension).
1212 Use this when EP degree is 1 and you want TP across experts without
1213 any all-to-all token dispatch. Typically combined with the standard
1214 :class:`~hyper_parallel.core.tensor_parallel.style.ColwiseParallel` /
1215 :class:`~hyper_parallel.core.tensor_parallel.style.RowwiseParallel`
1216 pattern for attention layers.
1218 Example::
1219 >>> tp_style = TensorParallel()
1220 >>> sharded_experts = tp_style.apply(experts_module, tp_device_mesh)
1221 """
1223 def apply(self, module: Module, device_mesh: DeviceMesh) -> Module:
1224 """Apply TP weight sharding to *module*.
1226 Args:
1227 module: A :class:`GroupedExperts` instance.
1228 device_mesh: 1-D TP device mesh (``mesh_dim_names=("tp",)``).
1230 Returns:
1231 The module with TP-sharded expert parameters.
1232 """
1233 return distribute_module(
1234 module,
1235 device_mesh,
1236 self._partition_fn,
1237 )
1239 @staticmethod
1240 def _partition_fn(name: str, module: Module, device_mesh: DeviceMesh) -> None:
1241 """Shard expert weights column-wise (w1/w3) or row-wise (w2).
1243 ``GroupedExperts`` weight layout is ``[num_experts, out_dim, in_dim]``
1244 so:
1246 - ``w1``/``w3``: shard ``Shard(1)`` → split ``hidden_dim``
1247 (column-wise analogue).
1248 - ``w2``: shard ``Shard(2)`` → split ``in_dim = hidden_dim``
1249 (row-wise analogue).
1251 Args:
1252 name: Submodule name (unused).
1253 module: The module whose parameters are being sharded.
1254 device_mesh: TP device mesh.
1255 """
1256 del name
1257 for key, param in _distribute_module_iter_params(module):
1258 if param is None:
1259 continue
1260 src = _distribute_module_param_source(param)
1261 requires_grad = bool(getattr(param, "requires_grad", True))
1262 # w1, w3: column-wise → Shard(1); w2: row-wise → Shard(2).
1263 shard_dim = 2 if key == "w2" else 1
1264 dt = distribute_tensor(src, device_mesh, [Shard(shard_dim)])
1265 new_param = _distribute_module_new_parameter(key, dt, requires_grad)
1266 _distribute_module_set_param(module, key, new_param)
1269# ---------------------------------------------------------------------------
1270# ExpertTensorParallel — combined EP + TP on a 2-D [ep, tp] mesh
1271# ---------------------------------------------------------------------------
1273class ExpertTensorParallel(ExpertParallel):
1274 """Combined Expert + Tensor Parallel on a 2-D ``[ep, tp]`` device mesh.
1276 Extends :class:`ExpertParallel` to operate on a 2-D mesh with named
1277 dimensions ``"ep"`` and ``"tp"``:
1279 - **Partition**: each expert weight ``[num_experts, out, in]`` is doubly
1280 sharded — ``Shard(0)`` along the EP dim (expert ownership) and
1281 ``Shard(1)``/``Shard(2)`` along the TP dim (column-wise / row-wise).
1282 - **Dispatch / Combine**: use only the 1-D ``device_mesh["ep"]`` sub-mesh
1283 so that token routing uses EP-group collectives, not the full 2-D mesh.
1285 Args:
1286 token_dispatcher: Token dispatch strategy. Supported values are
1287 ``"all_to_all"`` and ``"deredundency"``.
1288 async_combine: Forwarded to :class:`ExpertParallel`. When ``True``,
1289 the combine all-to-all is launched asynchronously for
1290 shared-expert overlap.
1292 Example::
1293 >>> etp_style = ExpertTensorParallel()
1294 >>> sharded = etp_style.apply(experts_module, ep_tp_2d_mesh)
1295 """
1297 def __init__(self, token_dispatcher: Union[str, bool] = "all_to_all", async_combine: bool = False) -> None:
1298 """Initialize ExpertTensorParallel.
1300 Args:
1301 async_combine: If ``True``, use asynchronous combine all-to-all.
1302 """
1303 super().__init__(token_dispatcher=token_dispatcher, async_combine=async_combine)
1305 def _dispatch_mesh(self, device_mesh: DeviceMesh) -> DeviceMesh:
1306 """Return the mesh used for token dispatch in ETP."""
1307 if self._token_dispatcher_name == "deredundency":
1308 raise NotImplementedError(
1309 "ExpertTensorParallel does not yet support "
1310 "token_dispatcher='deredundency'. Use ExpertParallel with a "
1311 "[oep, iep] mesh, or add [oep, iep, tp] mesh handling first."
1312 )
1313 return device_mesh["ep"]
1315 def _token_dispatch(self, module: Module, inputs, device_mesh: DeviceMesh):
1316 """Dispatch tokens using only the EP sub-mesh.
1318 Args:
1319 module: The ``GroupedExperts`` module.
1320 inputs: Tuple ``(routed_input, num_tokens_per_expert)`` or
1321 ``(routed_input, num_tokens_per_expert, routed_probs)``.
1322 device_mesh: 2-D device mesh with dims ``("ep", "tp")``.
1324 Returns:
1325 Tuple ``(permuted_local_input, local_token_counts)`` or
1326 ``(permuted_local_input, local_token_counts, permuted_probs)``
1327 depending on whether *routed_probs* was provided.
1328 """
1329 dispatch_mesh = self._dispatch_mesh(device_mesh)
1330 # Gate: _dispatch_mesh raises NotImplementedError for deredundency,
1331 # keeping dispatch/combine symmetric (both reject deredundency).
1332 dispatch_result = self._token_dispatcher.dispatch(module, inputs, dispatch_mesh)
1333 ctx = dispatch_result[-1]
1334 # Store context in module attribute for _token_combine to read.
1335 # Using module attribute ensures each module has its own context,
1336 # solving the instance sharing problem when the same ExpertParallel
1337 # style object is applied to multiple GroupedExperts modules.
1338 # pylint: disable=W0212
1339 # Store context in module attribute for _token_combine to read.
1340 module._ep_dispatch_ctx = ctx
1341 # dispatch_result is either (permuted, local_counts, ctx) or
1342 # (permuted, local_counts, permuted_probs, ctx); return all but ctx.
1343 return dispatch_result[:-1]
1345 def _token_combine(self, module: Module, routed_output, device_mesh: DeviceMesh):
1346 """Combine tokens using only the EP sub-mesh.
1348 When ``async_combine=True``, launches the combine all-to-all
1349 asynchronously via :meth:`self._token_dispatcher.combine_start`.
1351 Args:
1352 module: The ``GroupedExperts`` module.
1353 routed_output: Expert output tensor in expert-major order.
1354 device_mesh: 2-D device mesh with dims ``("ep", "tp")``.
1356 Returns:
1357 Token tensor in the original token-major layout.
1359 Raises:
1360 RuntimeError: If dispatch context is not found.
1361 """
1362 # pylint: disable=W0212
1363 # Read dispatch context from module attribute set by _token_dispatch.
1364 ctx = getattr(module, "_ep_dispatch_ctx", None)
1365 if ctx is None:
1366 raise RuntimeError(
1367 "_token_combine called but no dispatch context found in module. "
1368 "This indicates _token_dispatch was not called before _token_combine, "
1369 "or the context was already consumed by a previous combine call."
1370 )
1372 # Note: Do NOT delete the context here. In PyTorch, the tensors in ctx
1373 # are captured by autograd graph and don't need the attribute. But in
1374 # MindSpore PyNative mode, deleting the attribute may break backward.
1375 # The context will be overwritten on the next forward call.
1377 dispatch_mesh = self._dispatch_mesh(device_mesh)
1379 if self.async_combine:
1380 handle = self._token_dispatcher.combine_start(
1381 routed_output, dispatch_mesh, ctx
1382 )
1383 # pylint: disable=W0212
1384 module._ep_combine_handle = handle
1385 return handle.wait()
1387 return self._token_dispatcher.combine(
1388 module, routed_output, dispatch_mesh, ctx,
1389 )
1391 def _partition_fn(
1392 self, name: str, module: Module, device_mesh: DeviceMesh
1393 ) -> None:
1394 """Shard expert weights along both EP (dim 0) and TP (dim 1 or 2).
1396 Weight layout ``[num_experts, out_dim, in_dim]``:
1398 - ``w1``/``w3``: ``[Shard(0), Shard(1)]`` — EP shards experts,
1399 TP splits hidden_dim (column-wise).
1400 - ``w2``: ``[Shard(0), Shard(2)]`` — EP shards experts, TP splits
1401 the input dimension (row-wise).
1403 Args:
1404 name: Submodule name (unused).
1405 module: The module whose parameters are being sharded.
1406 device_mesh: 2-D device mesh with dims ``("ep", "tp")``.
1407 """
1408 del name
1409 for key, param in _distribute_module_iter_params(module):
1410 if param is None:
1411 continue
1412 src = _distribute_module_param_source(param)
1413 requires_grad = bool(getattr(param, "requires_grad", True))
1414 # EP shards expert ownership (dim 0); TP shards weight dim.
1415 tp_dim = 2 if key == "w2" else 1
1416 dt = distribute_tensor(src, device_mesh, [Shard(0), Shard(tp_dim)])
1417 new_param = _distribute_module_new_parameter(key, dt, requires_grad)
1418 _distribute_module_set_param(module, key, new_param)