Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / expert_parallel / expert_parallel.py: 93%
370 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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 # Re-exported for the UT, which applies it directly.
37 "_AsyncA2ALazyBwd",
38]
40import logging
41from abc import ABC, abstractmethod
42from dataclasses import dataclass
43from typing import Any, List, Optional, Tuple, Union
45import torch
46from torch.nn import Module
48from hyper_parallel.core.expert_parallel.static_splits import (
49 get_static_plan,
50 static_plan_key,
51 store_static_plan,
52)
53from hyper_parallel.core.dtensor.device_mesh import DeviceMesh
54from hyper_parallel.core.dtensor.dtensor import (
55 distribute_module,
56 distribute_tensor,
57 _distribute_module_iter_params,
58 _distribute_module_new_parameter,
59 _distribute_module_param_source,
60 _distribute_module_set_param,
61)
62from hyper_parallel.core.dtensor.placement_types import Shard
63from hyper_parallel.core.tensor_parallel.style import ParallelStyle
64from hyper_parallel.core.utils.communication import (
65 _AsyncA2ALazyBwd,
66 differentiable_all_gather_concat,
67 differentiable_all_to_all_single,
68 differentiable_all_to_all_single_async,
69 differentiable_reduce_scatter,
70 exchange_splits_via_all_to_all,
71 gather_counts_via_all_gather,
72 wait_async_tensor,
73)
76class AsyncHandle:
77 """Idempotent wait handle for an async collective operation.
79 Wraps the async tensor returned by
80 :meth:`Platform.differentiable_all_to_all_single_async` and provides a
81 :meth:`wait` method that is safe to call multiple times.
82 """
84 def __init__(self, async_tensor) -> None:
85 self._tensor = async_tensor
86 self._waited = False
88 def wait(self):
89 """Wait for the async collective to complete.
91 Idempotent — the first call blocks until the collective finishes;
92 subsequent calls are no-ops.
94 Returns:
95 The now-materialised result tensor.
96 """
97 if not self._waited:
98 wait_async_tensor(self._tensor)
99 self._waited = True
100 return self._tensor
103# ---------------------------------------------------------------------------
104# Token permutation helpers
105# ---------------------------------------------------------------------------
107def _generate_permute_indices(
108 tokens_per_expert_group,
109 experts_per_rank: int,
110 num_ranks: int,
111):
112 """Generate permutation indices for rank-major → expert-major reordering.
114 After all-to-all, received tokens are laid out in rank-major order::
116 [rank0·expert0 tokens | rank0·expert1 tokens | ... |
117 rank1·expert0 tokens | rank1·expert1 tokens | ...]
119 Expert computation requires expert-major order::
121 [all tokens for local expert 0 | all tokens for local expert 1 | ...]
123 Args:
124 tokens_per_expert_group: 1-D integer tensor of shape
125 ``[num_ranks * experts_per_rank]``. Entry ``[r * E + e]`` is the
126 number of tokens received from rank ``r`` for local expert ``e``.
127 experts_per_rank: Number of experts owned by each rank.
128 num_ranks: EP degree (total number of ranks in the EP group).
130 Returns:
131 Tuple of:
133 - ``permuted_indices``: 1-D long tensor of length
134 ``total_received_tokens``. ``permuted_indices[i]`` is the source
135 position in the rank-major buffer for destination position ``i`` in
136 the expert-major buffer.
137 - ``num_tokens_per_expert``: 1-D integer tensor of length
138 ``experts_per_rank`` with the token count per local expert.
139 """
140 counts = tokens_per_expert_group # [num_ranks * experts_per_rank]
142 # num_tokens_per_expert[e] = Σ_r counts[r * E + e]
143 counts_2d = counts.view(num_ranks, experts_per_rank) # [R, E]
144 num_tokens_per_expert = counts_2d.sum(dim=0) # [E]
146 # ``total`` must be a host int because ``arange`` needs a scalar size.
147 # That single D2H drain is unavoidable. Everything else stays on
148 # device — no per-block ``.item()`` in a loop.
149 total = int(num_tokens_per_expert.sum())
150 if total == 0:
151 return counts.new_zeros(0, dtype=counts.dtype), num_tokens_per_expert
153 # ---- Vectorized expert-major permutation, no host stalls -----------
154 # Source offsets in the rank-major receive buffer for each (r, e) block.
155 src_offsets_rm = counts.cumsum(0) - counts # [R*E], starts of each block
156 # Reorder src offsets to expert-major iteration order: block (e, r).
157 src_offsets_em = (
158 src_offsets_rm.view(num_ranks, experts_per_rank).T.contiguous().view(-1)
159 ) # [E*R]
160 # Counts in expert-major iteration order.
161 counts_em = counts_2d.T.contiguous().view(-1) # [E*R]
163 # ``repeat_interleave`` expands each block's src start to one entry per
164 # token in that block — gives the source position of each output token.
165 block_src_starts = src_offsets_em.repeat_interleave(counts_em) # [total]
167 # Destination block starts in expert-major order, then expanded. The
168 # ``arange(total) - dst_block_starts_per_token`` produces 0..n-1 within
169 # each block, i.e. the intra-block offset.
170 dst_block_starts = counts_em.cumsum(0) - counts_em # [E*R]
171 dst_block_starts_per_token = dst_block_starts.repeat_interleave(counts_em)
172 intra = torch.arange(0, total, device=counts.device) - dst_block_starts_per_token
174 permuted_indices = (block_src_starts + intra).long()
175 return permuted_indices, num_tokens_per_expert
178def _permute(x, tokens_per_expert_group, ep_degree: int, num_local_experts: int):
179 """Apply rank-major → expert-major permutation to routed tokens.
181 Args:
182 x: Received token tensor of shape
183 ``[sum(tokens_per_expert_group), *feature_dims]``.
184 tokens_per_expert_group: 1-D integer tensor of shape
185 ``[ep_degree * num_local_experts]`` (output of the first
186 all-to-all that exchanges token counts).
187 ep_degree: EP group size (number of ranks).
188 num_local_experts: Number of experts owned by this rank.
190 Returns:
191 Tuple of:
193 - ``original_shape``: shape of *x* before permutation.
194 - ``permuted_x``: tokens reordered to expert-major layout.
195 - ``permuted_indices``: permutation indices (needed for
196 :func:`_unpermute`).
197 - ``num_tokens_per_expert``: token count per local expert.
198 """
199 original_shape = x.shape
200 permuted_indices, num_tokens_per_expert = _generate_permute_indices(
201 tokens_per_expert_group, num_local_experts, ep_degree
202 )
203 # ``x[permuted_indices]`` works for empty indices too (returns a
204 # shape-0 tensor with a real grad_fn). Avoid the early-return with
205 # ``new_zeros`` which would produce a leaf tensor without grad_fn and
206 # silently break autograd for ranks that happen to receive zero tokens.
207 permuted_x = x[permuted_indices]
208 return original_shape, permuted_x, permuted_indices, num_tokens_per_expert
211def _unpermute(out, original_shape, permuted_indices):
212 """Reverse the permutation applied by :func:`_permute`.
214 Args:
215 out: Expert-major output tensor of shape
216 ``[sum(num_tokens_per_expert), *feature_dims]``.
217 original_shape: Shape before permutation (from :func:`_permute`).
218 permuted_indices: Permutation indices from :func:`_permute`.
220 Returns:
221 Token tensor restored to the rank-major layout received after
222 all-to-all, with shape ``original_shape``.
223 """
224 # ``result[permuted_indices] = out`` is a differentiable scatter that
225 # also handles the empty-index case (no-op assignment, but autograd
226 # still connects ``result`` back to ``out``). Do NOT short-circuit
227 # with a bare ``new_zeros`` — that returns a leaf tensor without
228 # grad_fn and the downstream combine a2a loses its backward path,
229 # which manifests as "element 0 of tensors does not require grad".
230 result = out.new_zeros(*original_shape)
231 result[permuted_indices] = out
232 return result
235# ---------------------------------------------------------------------------
236# DispatchContext — state shared between token dispatch and combine
237# ---------------------------------------------------------------------------
239@dataclass
240class DispatchContext:
241 """Intermediate state between token dispatch and combine.
243 Stored in ``module._ep_dispatch_ctx`` for a single forward pass.
244 This solves the instance sharing problem when the same ExpertParallel
245 style object is applied to multiple layers:
247 Example problem (before this fix):
248 ep_style = ExpertParallel()
249 ep_style.apply(layer1.experts, mesh) # registers hooks
250 ep_style.apply(layer2.experts, mesh) # reuses same ep_style
252 # During forward:
253 # layer1.dispatch writes to ep_style._state_stack
254 # layer2.dispatch pushes to same stack ← INTERLEAVING
255 # layer1.combine pops wrong state (LIFO violation)
257 Solution: Store context per-module, not per-style-instance.
259 Built by :meth:`AllToAllTokenDispatcher.dispatch` and consumed by
260 :meth:`AllToAllTokenDispatcher.combine`. The caller
261 (e.g. :class:`ExpertParallel`) stores this on the module between the
262 paired dispatch/combine calls.
263 """
265 input_splits: List[int]
266 output_splits: List[int]
267 input_shape: Tuple[int, ...]
268 permuted_indices: Any
269 probs_input_shape: Optional[Tuple[int, ...]] = None
272@dataclass
273class DeredundencyDispatchContext(DispatchContext):
274 """State shared by deredundency dispatch and combine.
276 The inherited split and permutation fields describe the inner-EP
277 all-to-all. The extra fields describe the OEP shared token view and the
278 whiteboard scatter used before the final reduce-scatter combine.
279 """
281 dispatch_indices: Optional[object] = None
282 router_coeff: Optional[object] = None
283 gathered_shape: Optional[tuple] = None
284 oep_size: int = 1
288@dataclass(frozen=True)
289class _DeredundencyMeshInfo:
290 """Resolved mesh metadata for two-stage deredundency token exchange."""
292 oep_group: object
293 iep_group: object
294 oep_size: int
295 iep_size: int
296 outer_rank: int
297 inner_rank: int
300def _get_deredundency_mesh_info(device_mesh: DeviceMesh) -> _DeredundencyMeshInfo:
301 """Resolve ``oep`` / ``iep`` groups from a 1-D or 2-D EP mesh."""
302 ndim = getattr(device_mesh, "ndim", 1)
303 if not isinstance(ndim, int):
304 ndim = 1
305 if ndim == 1:
306 return _DeredundencyMeshInfo(
307 oep_group=None,
308 iep_group=device_mesh.get_group(),
309 oep_size=1,
310 iep_size=device_mesh.size(),
311 outer_rank=0,
312 inner_rank=device_mesh.get_local_rank(),
313 )
314 if ndim != 2:
315 raise ValueError(
316 "DeredundencyTokenDispatcher expects a 1-D EP mesh or a 2-D "
317 f"[oep, iep] EP mesh, but got ndim={ndim}."
318 )
320 mesh_dim_names = getattr(device_mesh, "mesh_dim_names", None) or ()
321 oep_dim = mesh_dim_names.index("oep") if "oep" in mesh_dim_names else 0
322 iep_dim = mesh_dim_names.index("iep") if "iep" in mesh_dim_names else 1
323 if oep_dim == iep_dim:
324 raise ValueError("DeredundencyTokenDispatcher requires distinct oep and iep mesh dimensions.")
326 return _DeredundencyMeshInfo(
327 oep_group=device_mesh.get_group(oep_dim),
328 iep_group=device_mesh.get_group(iep_dim),
329 oep_size=device_mesh.size(oep_dim),
330 iep_size=device_mesh.size(iep_dim),
331 outer_rank=device_mesh.get_local_rank(oep_dim),
332 inner_rank=device_mesh.get_local_rank(iep_dim),
333 )
336def _expert_offsets_by_source(tokens_per_expert_by_source):
337 """Return absolute expert-block offsets for every OEP source rank."""
338 source_totals = tokens_per_expert_by_source.sum(dim=1)
339 source_offsets = source_totals.cumsum(0) - source_totals
340 return (
341 tokens_per_expert_by_source.cumsum(dim=1)
342 - tokens_per_expert_by_source
343 + source_offsets.view(tokens_per_expert_by_source.shape[0], 1)
344 )
347def _expand_dispatch_blocks(block_counts, block_starts, device):
348 """Expand contiguous expert blocks into gather-view token indices."""
349 total = int(block_counts.sum())
350 if total == 0:
351 return block_counts.new_zeros(0, dtype=block_counts.dtype).long()
352 repeated_starts = block_starts.repeat_interleave(block_counts)
353 block_offsets = block_counts.cumsum(0) - block_counts
354 repeated_offsets = block_offsets.repeat_interleave(block_counts)
355 intra_block_offsets = torch.arange(0, total, device=device) - repeated_offsets
356 return (repeated_starts + intra_block_offsets).long()
359def _generate_deredundency_dispatch_indices(
360 tokens_per_expert_by_source,
361 expert_start: int,
362 iep_size: int,
363 num_local_experts: int,
364):
365 """Generate gather-view indices ordered by IEP destination rank.
367 ``tokens_per_expert_by_source`` is shaped ``[oep_size, num_experts]`` and
368 describes each source rank's expert-major routed buffer after the OEP
369 all-gather. The returned indices select the current outer expert range
370 and order it as ``[iep_dst, local_expert, oep_source]`` so each IEP
371 destination chunk keeps local-expert blocks contiguous for the later
372 rank-major → expert-major permutation.
373 """
374 oep_size = tokens_per_expert_by_source.shape[0]
375 expert_end = expert_start + iep_size * num_local_experts
376 expert_offsets = _expert_offsets_by_source(tokens_per_expert_by_source)
377 selected_counts = tokens_per_expert_by_source[:, expert_start:expert_end].view(
378 oep_size, iep_size, num_local_experts,
379 )
380 selected_offsets = expert_offsets[:, expert_start:expert_end].view(
381 oep_size, iep_size, num_local_experts,
382 )
383 counts_by_destination = selected_counts.permute(1, 2, 0).contiguous()
384 offsets_by_destination = selected_offsets.permute(1, 2, 0).contiguous()
386 block_counts = counts_by_destination.view(-1)
387 token_counts_by_destination_expert = selected_counts.sum(dim=0).contiguous().view(-1)
388 block_starts = offsets_by_destination.view(-1)
389 dispatch_indices = _expand_dispatch_blocks(
390 block_counts, block_starts, tokens_per_expert_by_source.device
391 )
392 return dispatch_indices, token_counts_by_destination_expert
395def _scale_by_router_coeff(tokens, router_coeff):
396 """Scale routed expert outputs by optional router coefficients."""
397 if router_coeff is None:
398 return tokens
399 if router_coeff.shape[0] != tokens.shape[0]:
400 raise ValueError(
401 "router_coeff length must match routed token count, got "
402 f"{router_coeff.shape[0]} and {tokens.shape[0]}."
403 )
404 coeff = router_coeff
405 if len(coeff.shape) == 1 and len(tokens.shape) > 1:
406 coeff = coeff.reshape((-1,) + (1,) * (len(tokens.shape) - 1))
407 return tokens * coeff
410def _scatter_add_first_dim(src, indices, output_shape):
411 """Scatter-add rows of ``src`` into a zero tensor along dim 0."""
412 result = src.new_zeros(*output_shape)
413 if len(src.shape) == 1:
414 scatter_indices = indices
415 else:
416 scatter_indices = indices.reshape((-1,) + (1,) * (len(src.shape) - 1)).expand(
417 -1, *src.shape[1:],
418 )
419 if hasattr(result, "scatter_add"):
420 return result.scatter_add(0, scatter_indices, src)
421 if hasattr(result, "index_add"):
422 return result.index_add(0, indices, src)
423 raise RuntimeError(
424 "DeredundencyTokenDispatcher.combine requires tensor scatter_add or "
425 "index_add support for exdispatch_idx accumulation."
426 )
429class _DeredundencyCombineHandle(AsyncHandle):
430 """Async handle that finishes deredundency combine post-processing."""
432 def __init__(
433 self,
434 async_tensor: object,
435 mesh_info: _DeredundencyMeshInfo,
436 ctx: DeredundencyDispatchContext,
437 ) -> None:
438 super().__init__(async_tensor)
439 self._mesh_info = mesh_info
440 self._ctx = ctx
441 self._combined: Optional[object] = None
443 def wait(self) -> object:
444 """Wait for IEP a2a, then finish OEP scatter/reduce combine once."""
445 if self._combined is None:
446 outer_output = super().wait()
447 weighted_output = _scale_by_router_coeff(outer_output, self._ctx.router_coeff)
448 combine_whiteboard = _scatter_add_first_dim(
449 weighted_output,
450 self._ctx.dispatch_indices,
451 self._ctx.gathered_shape,
452 )
453 if self._ctx.oep_size == 1:
454 self._combined = combine_whiteboard
455 else:
456 self._combined = differentiable_reduce_scatter(
457 combine_whiteboard,
458 self._ctx.oep_size,
459 0,
460 "sum",
461 self._mesh_info.oep_group,
462 )
463 return self._combined
466# ---------------------------------------------------------------------------
467# BaseExpertParallel — abstract base for all-to-all EP strategies
468# ---------------------------------------------------------------------------
470class BaseExpertParallel(ParallelStyle, ABC):
471 """Abstract base class for Expert Parallel strategies with token dispatch.
473 Subclasses implement :meth:`_partition_fn`, :meth:`_token_dispatch`, and
474 :meth:`_token_combine`; this class wires them into :func:`distribute_module`.
475 """
477 def apply(self, module: Module, device_mesh: DeviceMesh) -> Module:
478 """Apply EP sharding and dispatch/combine hooks to *module*.
480 Args:
481 module: A :class:`~hyper_parallel.core.expert_parallel.moe.GroupedExperts`
482 instance to shard.
483 device_mesh: Device mesh for this EP strategy.
485 Returns:
486 The module with distributed parameters and dispatch/combine hooks.
487 """
488 return distribute_module(
489 module,
490 device_mesh,
491 self._partition_fn,
492 self._token_dispatch,
493 self._token_combine,
494 )
496 @abstractmethod
497 def _partition_fn(
498 self, name: str, module: Module, device_mesh: DeviceMesh
499 ) -> None:
500 """Shard module parameters according to this strategy.
502 Args:
503 name: Submodule name.
504 module: The module whose parameters are being sharded.
505 device_mesh: Device mesh for this EP strategy.
506 """
508 @abstractmethod
509 def _token_dispatch(self, module: Module, inputs, device_mesh: DeviceMesh):
510 """Pre-hook: route input tokens to their assigned ranks.
512 Args:
513 module: The ``GroupedExperts`` module.
514 inputs: Forward inputs tuple.
515 device_mesh: Device mesh for this EP strategy.
517 Returns:
518 Transformed inputs for local expert computation.
519 """
521 @abstractmethod
522 def _token_combine(self, module: Module, routed_output, device_mesh: DeviceMesh):
523 """Post-hook: gather expert outputs back to the originating ranks.
525 Args:
526 module: The ``GroupedExperts`` module.
527 routed_output: Expert output tensor in expert-major order.
528 device_mesh: Device mesh for this EP strategy.
530 Returns:
531 Token tensor in the original token-major layout.
532 """
535# ---------------------------------------------------------------------------
536# AllToAllTokenDispatcher — token dispatch/combine via all-to-all
537# ---------------------------------------------------------------------------
539class AllToAllTokenDispatcher:
540 """Token dispatch and combine via all-to-all for expert parallelism.
542 Provides :meth:`dispatch` and :meth:`combine` as static methods that
543 receive and return a :class:`DispatchContext` object. This decouples
544 the all-to-all token routing logic from the parallel style class so
545 that it can be reused or tested independently.
547 Callers (e.g. :class:`ExpertParallel`) are responsible for storing the
548 context between the paired dispatch/combine calls.
549 """
551 @staticmethod
552 def dispatch(module: Module, inputs: tuple, device_mesh: DeviceMesh) -> tuple:
553 """Dispatch tokens to their assigned ranks via all-to-all.
555 Called as an ``input_fn`` hook by :func:`distribute_module`. Receives
556 the module's forward inputs and returns transformed inputs.
558 Args:
559 module: The ``GroupedExperts`` module.
560 inputs: Tuple ``(routed_input, num_tokens_per_expert)`` or
561 ``(routed_input, num_tokens_per_expert, permuted_probs)`` where
562 ``routed_input`` has shape ``[total_tokens, dim]``, and
563 ``num_tokens_per_expert`` has shape ``[num_experts]``, and
564 ``permuted_probs`` has shape ``[total_tokens]``.
565 device_mesh: EP device mesh (1-D).
567 Returns:
568 Tuple ``(permuted_local_input, local_token_counts, ctx)`` or
569 ``(permuted_local_input, local_token_counts, permuted_probs, ctx)``
570 depending on whether *permuted_probs* was provided -
571 the first two elements are the transformed inputs for local
572 expert computation; *ctx* is a :class:`DispatchContext`
573 carrying the updated state to be stored by the caller.
574 """
575 del module # module unused, kept for API consistency
576 routed_input, num_tokens_per_expert = inputs[0], inputs[1]
577 permuted_probs = inputs[2] if len(inputs) > 2 else None
578 ep_group = device_mesh.get_group()
579 ep_size = device_mesh.size()
580 num_local_experts = num_tokens_per_expert.shape[0] // ep_size
582 # --- Steps 1-2: token counts and the split vectors ---
583 # Under a static routing plan the answer never changes, so it is computed once and
584 # reused: that is what removes the per-layer counts a2a and its two D2H syncs.
585 plan_key = static_plan_key(
586 num_tokens_per_expert.shape[0], ep_size, num_local_experts,
587 device=num_tokens_per_expert.device,
588 )
589 cached_plan = get_static_plan(plan_key)
590 if cached_plan is not None:
591 counts_out, input_splits, output_splits = cached_plan
592 else:
593 # The counts exchange is shared with the other dispatchers (see
594 # ``core.utils.communication``): async with an explicit wait, because the
595 # implicit cross-stream sync is NCCL-only and on HCCL the compute stream
596 # could read ``counts_out`` before the collective write is visible.
597 counts_out = exchange_splits_via_all_to_all(num_tokens_per_expert, ep_group)
598 # counts_out shape: [ep_size * num_local_experts]
599 # counts_out[r * num_local_experts + e] = tokens from rank r for expert e
601 # input_splits[r] = tokens this rank sends to rank r
602 # output_splits[r] = tokens this rank receives from rank r
603 # Reshape to [ep_size, num_local_experts] and sum per rank on device;
604 # a single ``tolist()`` drains the rank-sum vector to host, replacing
605 # ``2 * ep_size`` scalar ``int()`` D2H syncs with 2.
606 input_splits = num_tokens_per_expert.view(ep_size, num_local_experts).sum(dim=1).tolist()
607 output_splits = counts_out.view(ep_size, num_local_experts).sum(dim=1).tolist()
608 store_static_plan(
609 plan_key, (counts_out, input_splits, output_splits),
610 f"(ep={ep_size}, local_experts={num_local_experts}, "
611 f"tokens/peer={input_splits[0] if input_splits else 0})",
612 )
614 # --- Step 3a: exchange actual tokens (differentiable) ---
615 dispatched = differentiable_all_to_all_single(
616 routed_input, input_splits, output_splits, group=ep_group,
617 )
619 # --- Step 4a: rank-major → expert-major permutation ---
620 input_shape, permuted, permuted_indices, local_counts = _permute(
621 dispatched, counts_out, ep_size, num_local_experts
622 )
624 # Steps 3b + 4b: exchange and reorder permuted_probs (when provided)
625 # alongside the tokens, keeping weights aligned with the expert-major
626 # token order. Extracted so dispatch stays under the Lizard NLOC cap.
627 probs_input_shape, permuted_probs_reordered = (
628 AllToAllTokenDispatcher._exchange_and_reorder_probs(
629 permuted_probs, input_splits, output_splits,
630 counts_out, ep_size, num_local_experts, ep_group,
631 )
632 )
634 # Build dispatch context for combine step.
635 # Caller (e.g., ExpertParallel._token_dispatch) is responsible for storing
636 # this context and passing it to combine(). This decouples dispatch/combine
637 # from module state and solves the instance sharing problem.
638 ctx = DispatchContext(
639 input_splits=input_splits,
640 output_splits=output_splits,
641 input_shape=input_shape,
642 permuted_indices=permuted_indices,
643 probs_input_shape=probs_input_shape,
644 )
646 if permuted_probs is not None:
647 return permuted, local_counts, permuted_probs_reordered, ctx
648 return permuted, local_counts, ctx
650 @staticmethod
651 def _exchange_and_reorder_probs(permuted_probs, input_splits, output_splits,
652 counts_out, ep_size, num_local_experts, ep_group):
653 """Exchange ``permuted_probs`` via a2a and reorder to expert-major.
655 Returns ``(probs_input_shape, permuted_probs_reordered)``; both are
656 ``None`` when *permuted_probs* is ``None`` (no-op, no second a2a).
657 """
658 if permuted_probs is None:
659 return None, None
660 # --- Step 3b: exchange permuted_probs via all-to-all (differentiable) ---
661 dispatched_probs = differentiable_all_to_all_single(
662 permuted_probs, input_splits, output_splits, group=ep_group,
663 )
664 # --- Step 4b: rank-major → expert-major permutation for probs ---
665 probs_input_shape, permuted_probs_reordered, _, _ = _permute(
666 dispatched_probs, counts_out, ep_size, num_local_experts
667 )
668 return probs_input_shape, permuted_probs_reordered
670 @staticmethod
671 def combine(module: Module, routed_output: object, device_mesh: DeviceMesh, ctx: DispatchContext) -> object:
672 """Gather expert outputs back to the originating ranks via all-to-all.
674 Called as an ``output_fn`` hook by :func:`distribute_module`.
675 Receives dispatch context from the caller (previously returned by dispatch).
677 Args:
678 module: The ``GroupedExperts`` module (unused, for API consistency).
679 routed_output: Expert output tensor in expert-major order,
680 shape ``[sum(local_counts), dim]``.
681 device_mesh: EP device mesh (1-D).
682 ctx: :class:`DispatchContext` previously returned by
683 :meth:`dispatch`.
685 Returns:
686 Token tensor in the original token-major layout,
687 shape ``[sum(input_splits), dim]``.
688 """
689 del module # module not used, kept for API consistency
690 ep_group = device_mesh.get_group()
692 # expert-major → rank-major
693 unpermuted = _unpermute(routed_output, ctx.input_shape, ctx.permuted_indices)
695 # reverse all-to-all (output/input splits are swapped)
696 combined = differentiable_all_to_all_single(
697 unpermuted,
698 ctx.output_splits, # was output, now becomes input
699 ctx.input_splits, # was input, now becomes output
700 group=ep_group,
701 )
702 return combined
704 @staticmethod
705 def combine_start(routed_output, device_mesh, ctx):
706 """Launch async combine all-to-all without waiting for completion.
708 Splits the combine into two phases so that the caller can overlap
709 the a2a communication with independent computation (e.g. a shared
710 expert forward pass). The caller must later call
711 :meth:`combine_wait` or ``handle.wait()`` to obtain the final
712 result.
714 Step 1 (synchronous, local): expert-major → rank-major unpermute.
715 Step 2 (asynchronous, cross-rank): reverse all-to-all.
717 Args:
718 routed_output: Expert output tensor in expert-major order,
719 shape ``[sum(local_counts), dim]``.
720 device_mesh: EP device mesh (1-D).
721 ctx: :class:`DispatchContext` previously returned by
722 :meth:`dispatch`.
724 Returns:
725 :class:`AsyncHandle` carrying the state needed by
726 :meth:`combine_wait`.
727 """
728 ep_group = device_mesh.get_group()
730 # expert-major → rank-major (local, no communication)
731 unpermuted = _unpermute(routed_output, ctx.input_shape, ctx.permuted_indices)
733 # async reverse all-to-all (output/input splits are swapped)
734 combined_async = differentiable_all_to_all_single_async(
735 unpermuted,
736 ctx.output_splits,
737 ctx.input_splits,
738 group=ep_group,
739 )
741 return AsyncHandle(combined_async)
743 @staticmethod
744 def combine_wait(handle):
745 """Wait for the async combine all-to-all to complete.
747 Args:
748 handle: :class:`AsyncHandle` returned by :meth:`combine_start`.
750 Returns:
751 Combined tensor in the original token-major layout.
752 """
753 return handle.wait()
756# ---------------------------------------------------------------------------
757# DeredundencyTokenDispatcher — token dispatch via OEP all-gather + IEP all-to-all
758# ---------------------------------------------------------------------------
760class DeredundencyTokenDispatcher:
761 """Token dispatch/combine via OEP all-gather plus IEP all-to-all.
763 This dispatcher keeps the same public contract as
764 :class:`AllToAllTokenDispatcher`, but decomposes the global EP all-to-all
765 into the deredundency flow described in
766 ``docs/moe_alltoall_deredundency_token_permutation.md``:
768 1. Form a shared token/count view across the OEP group.
769 2. Select only the current outer expert range.
770 3. Send selected tokens to concrete local-expert ranks inside the IEP
771 group.
772 4. Sort received tokens into local expert-major order.
774 For a 2-D mesh, dimension ``"oep"`` / ``0`` is the outer group and
775 ``"iep"`` / ``1`` is the inner group. A 1-D mesh is treated as
776 ``oep_size == 1`` and degenerates to the standard all-to-all data flow.
777 """
779 @staticmethod
780 def _oep_gather_for_dispatch(
781 num_tokens_per_expert,
782 routed_input,
783 router_coeff,
784 mesh_info: _DeredundencyMeshInfo,
785 ) -> tuple:
786 """All-gather token counts and routed input across the OEP group.
788 Args:
789 num_tokens_per_expert: Token count per expert ``[num_experts]``.
790 routed_input: Routed token tensor ``[total_tokens, dim]``.
791 router_coeff: Optional router coefficients ``[total_tokens]``.
792 mesh_info: Resolved OEP/IEP mesh descriptor.
794 Returns:
795 Tuple ``(gathered_counts, gathered_routed, gathered_router_coeff)``
796 where ``gathered_counts`` has shape ``[oep_size, num_experts]``,
797 ``gathered_routed`` has shape ``[oep_size * total_tokens, dim]``,
798 and ``gathered_router_coeff`` is the gathered coefficients or None.
800 Raises:
801 ValueError: If routed token counts differ across OEP ranks.
802 """
803 if mesh_info.oep_size == 1:
804 gathered_counts = num_tokens_per_expert.view(1, num_tokens_per_expert.shape[0])
805 return gathered_counts, routed_input, router_coeff
807 gathered_counts = gather_counts_via_all_gather(
808 num_tokens_per_expert, mesh_info.oep_group, mesh_info.oep_size,
809 )
810 source_token_totals = gathered_counts.sum(dim=1).tolist()
811 if any(total != routed_input.shape[0] for total in source_token_totals):
812 raise ValueError(
813 "DeredundencyTokenDispatcher requires equal routed token "
814 "counts within each OEP group because the shared token view "
815 f"uses all-gather, got totals {source_token_totals}."
816 )
817 gathered_routed = differentiable_all_gather_concat(
818 routed_input, mesh_info.oep_group, mesh_info.oep_size, 0,
819 )
820 if router_coeff is None:
821 gathered_router_coeff = None
822 else:
823 gathered_router_coeff = differentiable_all_gather_concat(
824 router_coeff, mesh_info.oep_group, mesh_info.oep_size, 0,
825 )
826 return gathered_counts, gathered_routed, gathered_router_coeff
828 @staticmethod
829 def dispatch(module: Module, inputs: tuple, device_mesh: DeviceMesh) -> tuple:
830 """Dispatch tokens using OEP all-gather and IEP all-to-all.
832 Args:
833 module: The ``GroupedExperts`` module (unused here).
834 inputs: Tuple ``(routed_input, num_tokens_per_expert)`` or
835 ``(routed_input, num_tokens_per_expert, router_coeff)`` where
836 ``routed_input`` has shape ``[total_tokens, dim]``,
837 ``num_tokens_per_expert`` has shape ``[num_experts]``, and
838 ``router_coeff`` is an optional 1-D tensor of shape
839 ``[total_tokens]``.
840 device_mesh: 1-D EP mesh or 2-D ``[oep, iep]`` EP mesh.
842 Returns:
843 Tuple ``(permuted_local_input, local_token_counts, ctx)`` with the
844 same meaning as :meth:`AllToAllTokenDispatcher.dispatch`. Unlike
845 :meth:`AllToAllTokenDispatcher.dispatch`, this always returns a
846 3-tuple — ``router_coeff`` (when provided) is consumed by the
847 combine step, not returned for in-expert weighting.
849 Raises:
850 ValueError: If the expert count is not divisible by the full EP
851 size represented by the deredundency mesh.
853 Note:
854 The third input is **router_coeff**, *not* ``permuted_probs``.
855 Both occupy ``inputs[2]`` with shape ``[total_tokens]`` but differ:
856 ``permuted_probs`` (:class:`AllToAllTokenDispatcher`) is exchanged
857 with tokens and applied **inside** experts (pre-w2 activation,
858 i.e. ``score_before_experts=False``); ``router_coeff`` (here) is
859 OEP-gathered, scattered with tokens, and applied in
860 :meth:`combine` to the expert **final output**. The w2/SiLU
861 nonlinearity makes the two mathematically distinct, so
862 deredundency does **not** support ``score_before_experts=False``
863 — passing scores as ``inputs[2]`` would be silently treated as
864 ``router_coeff``. Use ``score_before_experts=True`` until
865 deredundency's weighting is redesigned.
866 """
867 del module
868 routed_input, num_tokens_per_expert = inputs[0], inputs[1]
869 router_coeff = inputs[2] if len(inputs) > 2 else None
870 if router_coeff is not None and router_coeff.shape[0] != routed_input.shape[0]:
871 raise ValueError(
872 "router_coeff length must match routed_input token count, got "
873 f"{router_coeff.shape[0]} and {routed_input.shape[0]}."
874 )
875 mesh_info = _get_deredundency_mesh_info(device_mesh)
876 ep_size = mesh_info.oep_size * mesh_info.iep_size
877 if num_tokens_per_expert.shape[0] % ep_size != 0:
878 raise ValueError(
879 "num_tokens_per_expert length must be divisible by the full "
880 f"EP size {ep_size}, got {num_tokens_per_expert.shape[0]}."
881 )
882 num_local_experts = num_tokens_per_expert.shape[0] // ep_size
883 experts_per_outer = mesh_info.iep_size * num_local_experts
884 expert_start = mesh_info.outer_rank * experts_per_outer
886 gathered_counts, gathered_routed, gathered_router_coeff = (
887 DeredundencyTokenDispatcher._oep_gather_for_dispatch(
888 num_tokens_per_expert, routed_input, router_coeff, mesh_info,
889 )
890 )
892 dispatch_indices, node_counts_per_expert = _generate_deredundency_dispatch_indices(
893 gathered_counts,
894 expert_start,
895 mesh_info.iep_size,
896 num_local_experts,
897 )
898 iep_input_splits = node_counts_per_expert.view(
899 mesh_info.iep_size, num_local_experts).sum(dim=1).tolist()
900 iep_counts_out, iep_output_splits = (
901 DeredundencyTokenDispatcher._iep_exchange_counts(
902 node_counts_per_expert, mesh_info, num_local_experts,
903 )
904 )
905 outer_routed_input = gathered_routed[dispatch_indices]
906 outer_router_coeff = (
907 None if gathered_router_coeff is None else gathered_router_coeff[dispatch_indices]
908 )
909 dispatched = differentiable_all_to_all_single(
910 outer_routed_input, iep_input_splits, iep_output_splits,
911 group=mesh_info.iep_group,
912 )
913 input_shape, permuted, permuted_indices, local_counts = _permute(
914 dispatched, iep_counts_out, mesh_info.iep_size, num_local_experts,
915 )
916 ctx = DeredundencyDispatchContext(
917 input_splits=iep_input_splits,
918 output_splits=iep_output_splits,
919 input_shape=input_shape,
920 permuted_indices=permuted_indices,
921 dispatch_indices=dispatch_indices,
922 router_coeff=outer_router_coeff,
923 gathered_shape=gathered_routed.shape,
924 oep_size=mesh_info.oep_size,
925 )
926 return permuted, local_counts, ctx
928 @staticmethod
929 def _iep_exchange_counts(node_counts_per_expert, mesh_info, num_local_experts):
930 """IEP all-to-all of per-expert counts; return (counts_out, output_splits)."""
931 iep_counts_out = exchange_splits_via_all_to_all(node_counts_per_expert, mesh_info.iep_group)
932 iep_output_splits = iep_counts_out.view(
933 mesh_info.iep_size, num_local_experts).sum(dim=1).tolist()
934 return iep_counts_out, iep_output_splits
936 @staticmethod
937 def combine(module: Module, routed_output: object, device_mesh: DeviceMesh,
938 ctx: DeredundencyDispatchContext) -> object:
939 """Gather expert outputs back to the originating ranks.
941 Args:
942 module: The ``GroupedExperts`` module (unused).
943 routed_output: Expert output tensor in expert-major order.
944 device_mesh: 1-D EP mesh or 2-D ``[oep, iep]`` EP mesh.
945 ctx: Context returned by :meth:`dispatch`.
947 Returns:
948 Token tensor in the original source-rank routed order.
949 """
950 del module
951 mesh_info = _get_deredundency_mesh_info(device_mesh)
952 DeredundencyTokenDispatcher._validate_combine_mesh(mesh_info, ctx)
954 unpermuted = _unpermute(routed_output, ctx.input_shape, ctx.permuted_indices)
955 outer_output = differentiable_all_to_all_single(
956 unpermuted,
957 ctx.output_splits,
958 ctx.input_splits,
959 group=mesh_info.iep_group,
960 )
962 weighted_output = _scale_by_router_coeff(outer_output, ctx.router_coeff)
963 combine_whiteboard = _scatter_add_first_dim(
964 weighted_output, ctx.dispatch_indices, ctx.gathered_shape,
965 )
966 if ctx.oep_size == 1:
967 return combine_whiteboard
969 return differentiable_reduce_scatter(
970 combine_whiteboard,
971 ctx.oep_size,
972 0,
973 "sum",
974 mesh_info.oep_group,
975 )
977 @staticmethod
978 def _validate_combine_mesh(
979 mesh_info: _DeredundencyMeshInfo,
980 ctx: DeredundencyDispatchContext,
981 ) -> None:
982 """Validate that dispatch context and combine mesh are compatible."""
983 if mesh_info.oep_size != ctx.oep_size:
984 raise ValueError(
985 "DeredundencyTokenDispatcher.combine received a context for "
986 f"oep_size={ctx.oep_size}, but the mesh resolves to oep_size={mesh_info.oep_size}."
987 )
989 @staticmethod
990 def combine_start(
991 routed_output: object,
992 device_mesh: DeviceMesh,
993 ctx: DeredundencyDispatchContext,
994 ) -> AsyncHandle:
995 """Launch async IEP combine all-to-all and defer deredundency post-processing.
997 The local expert-major → rank-major unpermute is performed
998 synchronously. The reverse IEP all-to-all is launched asynchronously,
999 and :meth:`combine_wait` finishes router weighting, whiteboard
1000 scatter-add, and optional OEP reduce-scatter after the async output is
1001 materialised.
1003 Args:
1004 routed_output: Expert output tensor in expert-major order.
1005 device_mesh: 1-D EP mesh or 2-D ``[oep, iep]`` EP mesh.
1006 ctx: Context returned by :meth:`dispatch`.
1008 Returns:
1009 :class:`AsyncHandle` carrying the pending IEP a2a and deredundency
1010 combine state.
1011 """
1012 mesh_info = _get_deredundency_mesh_info(device_mesh)
1013 DeredundencyTokenDispatcher._validate_combine_mesh(mesh_info, ctx)
1015 unpermuted = _unpermute(routed_output, ctx.input_shape, ctx.permuted_indices)
1016 outer_output_async = differentiable_all_to_all_single_async(
1017 unpermuted,
1018 ctx.output_splits,
1019 ctx.input_splits,
1020 group=mesh_info.iep_group,
1021 )
1022 return _DeredundencyCombineHandle(outer_output_async, mesh_info, ctx)
1024 @staticmethod
1025 def combine_wait(handle: AsyncHandle) -> object:
1026 """Wait for async deredundency combine and return the final tensor.
1028 Args:
1029 handle: :class:`AsyncHandle` returned by :meth:`combine_start`.
1031 Returns:
1032 Token tensor in the original source-rank routed order.
1033 """
1034 return handle.wait()
1037logger = logging.getLogger(__name__)
1039_TOKEN_DISPATCHERS = {
1040 "all_to_all": AllToAllTokenDispatcher,
1041 "deredundency": DeredundencyTokenDispatcher,
1042}
1045def _resolve_token_dispatcher(token_dispatcher: str):
1046 """Resolve a token dispatcher name to its implementation class."""
1047 try:
1048 return _TOKEN_DISPATCHERS[token_dispatcher]
1049 except KeyError as exc:
1050 supported = "', '".join(sorted(_TOKEN_DISPATCHERS))
1051 raise ValueError(
1052 f"token_dispatcher must be one of '{supported}', got {token_dispatcher!r}."
1053 ) from exc
1056def _get_flattened_ep_mesh(device_mesh: DeviceMesh) -> DeviceMesh:
1057 """Return a 1-D EP mesh, flattening a 2-D deredundency mesh if needed."""
1058 if getattr(device_mesh, "ndim", 1) == 1:
1059 return device_mesh
1060 mesh_dim_names = getattr(device_mesh, "mesh_dim_names", None) or ()
1061 if "ep" in mesh_dim_names or "ep" in device_mesh.get_flatten_mapping():
1062 return device_mesh["ep"]
1063 if set(mesh_dim_names) == {"oep", "iep"}:
1064 return device_mesh.flatten("ep")
1065 raise ValueError(
1066 "Deredundency ExpertParallel expects a 1-D EP mesh or a 2-D "
1067 "[oep, iep] mesh when partitioning expert weights."
1068 )
1071# ---------------------------------------------------------------------------
1072# ExpertParallel — standard all-to-all EP
1073# ---------------------------------------------------------------------------
1075class ExpertParallel(BaseExpertParallel):
1076 """Expert Parallel: shard experts across ranks via all-to-all token routing.
1078 Applies :meth:`apply` to a :class:`GroupedExperts` module:
1080 1. **Partition** — distributes expert weights on dim 0 (``Shard(0)``) so
1081 each rank holds ``num_experts // ep_degree`` local experts.
1082 2. **Token dispatch** (forward pre-hook) — two-step all-to-all:
1083 a. Exchange token counts (non-differentiable).
1084 b. Exchange actual tokens (differentiable, gradient flows back).
1085 Followed by rank-major → expert-major permutation.
1086 3. **Token combine** (forward post-hook) — expert-major → rank-major
1087 unpermute, then reverse all-to-all (differentiable).
1089 Token collectives use Torch autograd; token-count exchanges use
1090 ``torch.distributed.all_to_all_single`` with an explicit wait.
1092 The token dispatcher is selectable. ``"all_to_all"`` uses
1093 :class:`AllToAllTokenDispatcher`; ``"deredundency"`` uses
1094 :class:`DeredundencyTokenDispatcher`.
1096 Args:
1097 token_dispatcher: Token dispatch strategy. Supported values are
1098 ``"all_to_all"`` and ``"deredundency"``.
1099 async_combine: When ``True``, the combine all-to-all is launched
1100 asynchronously so that the caller (e.g. :class:`MoE`) can
1101 overlap it with shared-expert computation. When ``False``
1102 (default), combine is fully synchronous — no overlap, identical
1103 to the baseline.
1105 Example::
1106 >>> ep_style = ExpertParallel()
1107 >>> sharded_experts = ep_style.apply(experts_module, ep_device_mesh)
1108 >>> # With async combine for shared-expert overlap:
1109 >>> ep_style = ExpertParallel(async_combine=True)
1110 >>> sharded_experts = ep_style.apply(experts_module, ep_device_mesh)
1111 """
1113 def __init__(self, token_dispatcher: Union[str, bool] = "all_to_all", async_combine: bool = False) -> None:
1114 """Initialize ExpertParallel.
1116 Args:
1117 token_dispatcher: Token dispatch strategy. Supported values are
1118 ``"all_to_all"`` and ``"deredundency"``.
1119 async_combine: If ``True``, use asynchronous combine all-to-all
1120 to overlap communication with shared-expert computation.
1121 """
1122 if isinstance(token_dispatcher, bool):
1123 async_combine = token_dispatcher
1124 token_dispatcher = "all_to_all"
1125 self._dispatch_ctx: Optional[DispatchContext] = None
1126 self.async_combine = async_combine
1127 self._token_dispatcher_name = token_dispatcher
1128 self._token_dispatcher = _resolve_token_dispatcher(token_dispatcher)
1130 def _token_dispatch(self, module: Module, inputs, device_mesh: DeviceMesh):
1131 """Dispatch tokens to their assigned ranks via all-to-all.
1133 Delegates to the configured token dispatcher and stores the
1134 returned :class:`DispatchContext` on the instance for the matching
1135 :meth:`_token_combine` call.
1137 Args:
1138 module: The ``GroupedExperts`` module.
1139 inputs: Tuple ``(routed_input, num_tokens_per_expert)`` or
1140 ``(routed_input, num_tokens_per_expert, routed_probs)``.
1141 device_mesh: EP device mesh (1-D).
1143 Returns:
1144 Tuple ``(permuted_local_input, local_token_counts)`` or
1145 ``(permuted_local_input, local_token_counts, permuted_probs)``
1146 depending on whether *routed_probs* was provided.
1147 """
1148 # Delegate to the configured dispatcher (all_to_all or deredundency).
1149 # Hard-coding AllToAllTokenDispatcher here would mismatch _token_combine
1150 # (which uses self._token_dispatcher) and break deredundency.
1151 dispatch_result = self._token_dispatcher.dispatch(module, inputs, device_mesh)
1152 ctx = dispatch_result[-1]
1153 # Store context in module attribute for _token_combine to read.
1154 # Using module attribute ensures each module has its own context,
1155 # solving the instance sharing problem when the same ExpertParallel
1156 # style object is applied to multiple GroupedExperts modules.
1157 # pylint: disable=W0212
1158 module._ep_dispatch_ctx = ctx
1159 # dispatch_result is either (permuted, local_counts, ctx) or
1160 # (permuted, local_counts, permuted_probs, ctx); return all but ctx.
1161 return dispatch_result[:-1]
1163 def _token_combine(self, module: Module, routed_output, device_mesh: DeviceMesh):
1164 """Gather expert outputs back to the originating ranks via all-to-all.
1166 When ``async_combine=True``, launches the combine all-to-all
1167 asynchronously and returns an :class:`AsyncCollectiveTensor`. The
1168 actual device-side wait is deferred until the downstream consumer
1169 (e.g. MoE unpermutation) first reads the tensor, enabling overlap
1170 with shared-expert computation.
1172 When ``async_combine=False`` (default), uses the synchronous
1173 :meth:`AllToAllTokenDispatcher.combine` — identical to the baseline.
1175 Args:
1176 module: The ``GroupedExperts`` module.
1177 routed_output: Expert output tensor in expert-major order.
1178 device_mesh: EP device mesh (1-D).
1180 Returns:
1181 Token tensor in the original token-major layout. When
1182 ``async_combine=True``, this may be an async collective tensor
1183 whose values are not yet materialised.
1185 Raises:
1186 RuntimeError: If dispatch context is not found (dispatch was not called).
1187 """
1188 # Read dispatch context from module attribute set by _token_dispatch.
1189 # pylint: disable=W0212
1190 ctx = getattr(module, "_ep_dispatch_ctx", None)
1191 if ctx is None:
1192 raise RuntimeError(
1193 "_token_combine called but no dispatch context found in module. "
1194 "This indicates _token_dispatch was not called before _token_combine, "
1195 "or the context was already consumed by a previous combine call."
1196 )
1198 # Note: Do NOT delete the context here. The tensors in ctx are
1199 # captured by the autograd graph, so the attribute is not needed for
1200 # backward and is simply overwritten on the next forward call.
1202 if self.async_combine:
1203 handle = self._token_dispatcher.combine_start(
1204 routed_output, device_mesh, ctx
1205 )
1206 # Return the async tensor. The first non-view access by the
1207 # downstream consumer (e.g. MoE unpermutation) will trigger the
1208 # implicit wait, overlapping with shared_expert computation.
1209 return handle.wait()
1211 return self._token_dispatcher.combine(
1212 module, routed_output, device_mesh, ctx,
1213 )
1215 def _partition_mesh(self, device_mesh: DeviceMesh) -> DeviceMesh:
1216 """Return the mesh used to shard expert weights."""
1217 if self._token_dispatcher_name == "deredundency":
1218 return _get_flattened_ep_mesh(device_mesh)
1219 return device_mesh
1221 def _partition_fn(
1222 self, name: str, module: Module, device_mesh: DeviceMesh
1223 ) -> None:
1224 """Shard all expert parameters along dim 0 (expert dimension).
1226 Args:
1227 name: Submodule name (unused).
1228 module: The module whose parameters are being sharded.
1229 device_mesh: EP device mesh.
1230 """
1231 del name
1232 partition_mesh = self._partition_mesh(device_mesh)
1233 for key, param in _distribute_module_iter_params(module):
1234 if param is None:
1235 continue
1236 src = _distribute_module_param_source(param)
1237 requires_grad = bool(getattr(param, "requires_grad", True))
1238 dt = distribute_tensor(src, partition_mesh, [Shard(0)])
1239 new_param = _distribute_module_new_parameter(key, dt, requires_grad)
1240 _distribute_module_set_param(module, key, new_param)
1243# ---------------------------------------------------------------------------
1244# TensorParallel — TP-only weight sharding for experts (no token dispatch)
1245# ---------------------------------------------------------------------------
1246class TensorParallel(ParallelStyle):
1247 """Tensor Parallel for expert weights (no token dispatch).
1249 Shards the ``GroupedExperts`` weight tensors in the column/row-wise
1250 pattern used by standard TP:
1252 - ``w1`` / ``w3``: ``Shard(1)`` — column-wise (hidden_dim dimension).
1253 - ``w2``: ``Shard(2)`` — row-wise (output dim dimension).
1255 Use this when EP degree is 1 and you want TP across experts without
1256 any all-to-all token dispatch. Typically combined with the standard
1257 :class:`~hyper_parallel.core.tensor_parallel.style.ColwiseParallel` /
1258 :class:`~hyper_parallel.core.tensor_parallel.style.RowwiseParallel`
1259 pattern for attention layers.
1261 Example::
1262 >>> tp_style = TensorParallel()
1263 >>> sharded_experts = tp_style.apply(experts_module, tp_device_mesh)
1264 """
1266 def apply(self, module: Module, device_mesh: DeviceMesh) -> Module:
1267 """Apply TP weight sharding to *module*.
1269 Args:
1270 module: A :class:`GroupedExperts` instance.
1271 device_mesh: 1-D TP device mesh (``mesh_dim_names=("tp",)``).
1273 Returns:
1274 The module with TP-sharded expert parameters.
1275 """
1276 return distribute_module(
1277 module,
1278 device_mesh,
1279 self._partition_fn,
1280 )
1282 @staticmethod
1283 def _partition_fn(name: str, module: Module, device_mesh: DeviceMesh) -> None:
1284 """Shard expert weights column-wise (w1/w3) or row-wise (w2).
1286 ``GroupedExperts`` weight layout is ``[num_experts, out_dim, in_dim]``
1287 so:
1289 - ``w1``/``w3``: shard ``Shard(1)`` → split ``hidden_dim``
1290 (column-wise analogue).
1291 - ``w2``: shard ``Shard(2)`` → split ``in_dim = hidden_dim``
1292 (row-wise analogue).
1294 Args:
1295 name: Submodule name (unused).
1296 module: The module whose parameters are being sharded.
1297 device_mesh: TP device mesh.
1298 """
1299 del name
1300 for key, param in _distribute_module_iter_params(module):
1301 if param is None:
1302 continue
1303 src = _distribute_module_param_source(param)
1304 requires_grad = bool(getattr(param, "requires_grad", True))
1305 # w1, w3: column-wise → Shard(1); w2: row-wise → Shard(2).
1306 shard_dim = 2 if key == "w2" else 1
1307 dt = distribute_tensor(src, device_mesh, [Shard(shard_dim)])
1308 new_param = _distribute_module_new_parameter(key, dt, requires_grad)
1309 _distribute_module_set_param(module, key, new_param)
1312# ---------------------------------------------------------------------------
1313# ExpertTensorParallel — combined EP + TP on a 2-D [ep, tp] mesh
1314# ---------------------------------------------------------------------------
1316class ExpertTensorParallel(ExpertParallel):
1317 """Combined Expert + Tensor Parallel on a 2-D ``[ep, tp]`` device mesh.
1319 Extends :class:`ExpertParallel` to operate on a 2-D mesh with named
1320 dimensions ``"ep"`` and ``"tp"``:
1322 - **Partition**: each expert weight ``[num_experts, out, in]`` is doubly
1323 sharded — ``Shard(0)`` along the EP dim (expert ownership) and
1324 ``Shard(1)``/``Shard(2)`` along the TP dim (column-wise / row-wise).
1325 - **Dispatch / Combine**: use only the 1-D ``device_mesh["ep"]`` sub-mesh
1326 so that token routing uses EP-group collectives, not the full 2-D mesh.
1328 Args:
1329 token_dispatcher: Token dispatch strategy. Supported values are
1330 ``"all_to_all"`` and ``"deredundency"``.
1331 async_combine: Forwarded to :class:`ExpertParallel`. When ``True``,
1332 the combine all-to-all is launched asynchronously for
1333 shared-expert overlap.
1335 Example::
1336 >>> etp_style = ExpertTensorParallel()
1337 >>> sharded = etp_style.apply(experts_module, ep_tp_2d_mesh)
1338 """
1340 def __init__(self, token_dispatcher: Union[str, bool] = "all_to_all", async_combine: bool = False) -> None:
1341 """Initialize ExpertTensorParallel.
1343 Args:
1344 async_combine: If ``True``, use asynchronous combine all-to-all.
1345 """
1346 super().__init__(token_dispatcher=token_dispatcher, async_combine=async_combine)
1348 def _dispatch_mesh(self, device_mesh: DeviceMesh) -> DeviceMesh:
1349 """Return the mesh used for token dispatch in ETP."""
1350 if self._token_dispatcher_name == "deredundency":
1351 raise NotImplementedError(
1352 "ExpertTensorParallel does not yet support "
1353 "token_dispatcher='deredundency'. Use ExpertParallel with a "
1354 "[oep, iep] mesh, or add [oep, iep, tp] mesh handling first."
1355 )
1356 return device_mesh["ep"]
1358 def _token_dispatch(self, module: Module, inputs, device_mesh: DeviceMesh):
1359 """Dispatch tokens using only the EP sub-mesh.
1361 Args:
1362 module: The ``GroupedExperts`` module.
1363 inputs: Tuple ``(routed_input, num_tokens_per_expert)`` or
1364 ``(routed_input, num_tokens_per_expert, routed_probs)``.
1365 device_mesh: 2-D device mesh with dims ``("ep", "tp")``.
1367 Returns:
1368 Tuple ``(permuted_local_input, local_token_counts)`` or
1369 ``(permuted_local_input, local_token_counts, permuted_probs)``
1370 depending on whether *routed_probs* was provided.
1371 """
1372 dispatch_mesh = self._dispatch_mesh(device_mesh)
1373 # Gate: _dispatch_mesh raises NotImplementedError for deredundency,
1374 # keeping dispatch/combine symmetric (both reject deredundency).
1375 dispatch_result = self._token_dispatcher.dispatch(module, inputs, dispatch_mesh)
1376 ctx = dispatch_result[-1]
1377 # Store context in module attribute for _token_combine to read.
1378 # Using module attribute ensures each module has its own context,
1379 # solving the instance sharing problem when the same ExpertParallel
1380 # style object is applied to multiple GroupedExperts modules.
1381 # pylint: disable=W0212
1382 # Store context in module attribute for _token_combine to read.
1383 module._ep_dispatch_ctx = ctx
1384 # dispatch_result is either (permuted, local_counts, ctx) or
1385 # (permuted, local_counts, permuted_probs, ctx); return all but ctx.
1386 return dispatch_result[:-1]
1388 def _token_combine(self, module: Module, routed_output, device_mesh: DeviceMesh):
1389 """Combine tokens using only the EP sub-mesh.
1391 When ``async_combine=True``, launches the combine all-to-all
1392 asynchronously via :meth:`self._token_dispatcher.combine_start`.
1394 Args:
1395 module: The ``GroupedExperts`` module.
1396 routed_output: Expert output tensor in expert-major order.
1397 device_mesh: 2-D device mesh with dims ``("ep", "tp")``.
1399 Returns:
1400 Token tensor in the original token-major layout.
1402 Raises:
1403 RuntimeError: If dispatch context is not found.
1404 """
1405 # pylint: disable=W0212
1406 # Read dispatch context from module attribute set by _token_dispatch.
1407 ctx = getattr(module, "_ep_dispatch_ctx", None)
1408 if ctx is None:
1409 raise RuntimeError(
1410 "_token_combine called but no dispatch context found in module. "
1411 "This indicates _token_dispatch was not called before _token_combine, "
1412 "or the context was already consumed by a previous combine call."
1413 )
1415 # Note: Do NOT delete the context here. The tensors in ctx are
1416 # captured by the autograd graph, so the attribute is not needed for
1417 # backward and is simply overwritten on the next forward call.
1419 dispatch_mesh = self._dispatch_mesh(device_mesh)
1421 if self.async_combine:
1422 handle = self._token_dispatcher.combine_start(
1423 routed_output, dispatch_mesh, ctx
1424 )
1425 return handle.wait()
1427 return self._token_dispatcher.combine(
1428 module, routed_output, dispatch_mesh, ctx,
1429 )
1431 def _partition_fn(
1432 self, name: str, module: Module, device_mesh: DeviceMesh
1433 ) -> None:
1434 """Shard expert weights along both EP (dim 0) and TP (dim 1 or 2).
1436 Weight layout ``[num_experts, out_dim, in_dim]``:
1438 - ``w1``/``w3``: ``[Shard(0), Shard(1)]`` — EP shards experts,
1439 TP splits hidden_dim (column-wise).
1440 - ``w2``: ``[Shard(0), Shard(2)]`` — EP shards experts, TP splits
1441 the input dimension (row-wise).
1443 Args:
1444 name: Submodule name (unused).
1445 module: The module whose parameters are being sharded.
1446 device_mesh: 2-D device mesh with dims ``("ep", "tp")``.
1447 """
1448 del name
1449 for key, param in _distribute_module_iter_params(module):
1450 if param is None:
1451 continue
1452 src = _distribute_module_param_source(param)
1453 requires_grad = bool(getattr(param, "requires_grad", True))
1454 # EP shards expert ownership (dim 0); TP shards weight dim.
1455 tp_dim = 2 if key == "w2" else 1
1456 dt = distribute_tensor(src, device_mesh, [Shard(0), Shard(tp_dim)])
1457 new_param = _distribute_module_new_parameter(key, dt, requires_grad)
1458 _distribute_module_set_param(module, key, new_param)