[moe] add EPMoELayer with a dispatcher / experts data contract

.github/scripts/ci/run_colossalai_batch.sh

@@ -76,6 +76,7 @@ case "${BATCH}" in
             tests/test_config/test_load_config.py
             tests/test_infer/test_async_engine/test_request_tracer.py
             tests/test_lazy/test_chatglm_cpu.py
+            tests/test_moe/test_ep_moe_layer_cpu.py
             tests/test_numerics/test_numerics_core.py
             tests/test_optimizer/test_lr_scheduler.py
             tests/test_pipeline/test_pipeline_utils/test_t5_pipeline_utils.py

.github/scripts/ci/run_cpu_unit_tests.sh

@@ -14,6 +14,7 @@ TESTS=(
     tests/test_config/test_load_config.py
     tests/test_infer/test_async_engine/test_request_tracer.py
     tests/test_lazy/test_chatglm_cpu.py
+    tests/test_moe/test_ep_moe_layer_cpu.py
     tests/test_numerics/test_numerics_core.py
     tests/test_optimizer/test_lr_scheduler.py
     tests/test_pipeline/test_pipeline_utils/test_t5_pipeline_utils.py

colossalai/moe/__init__.py

@@ -1,26 +1,34 @@
-"""MoE operations, loaded on demand to keep package imports lightweight."""
+"""MoE operations and the expert-parallel layer, loaded on demand to keep package imports lightweight."""
 
 from importlib import import_module
 
-__all__ = [
-    "AllGather",
-    "AllToAll",
-    "EPGradScalerIn",
-    "EPGradScalerOut",
-    "HierarchicalAllToAll",
-    "MoeCombine",
-    "MoeDispatch",
-    "ReduceScatter",
-    "all_to_all_uneven",
-    "moe_cumsum",
-]
+_SUBMODULE = {
+    "AllGather": "_operation",
+    "AllToAll": "_operation",
+    "EPGradScalerIn": "_operation",
+    "EPGradScalerOut": "_operation",
+    "HierarchicalAllToAll": "_operation",
+    "MoeCombine": "_operation",
+    "MoeDispatch": "_operation",
+    "ReduceScatter": "_operation",
+    "all_to_all_uneven": "_operation",
+    "moe_cumsum": "_operation",
+    "DispatchOutput": "dispatcher",
+    "Dispatcher": "dispatcher",
+    "NCCLAllToAllDispatcher": "dispatcher",
+    "Experts": "experts",
+    "GroupedMMExperts": "experts",
+    "ReferenceExperts": "experts",
+    "EPMoELayer": "layer",
+}
+__all__ = list(_SUBMODULE)
 
 
 def __getattr__(name):
-    if name in __all__:
-        operation = getattr(import_module("._operation", __name__), name)
-        globals()[name] = operation
-        return operation
+    if name in _SUBMODULE:
+        value = getattr(import_module(f".{_SUBMODULE[name]}", __name__), name)
+        globals()[name] = value
+        return value
     raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
 
 

colossalai/moe/dispatcher.py

@@ -0,0 +1,115 @@
+"""Token dispatchers for expert parallelism and the data contract they hand to the experts.
+
+A dispatcher moves each token's copies to the ranks that hold its top-k experts (``dispatch``) and brings the expert
+outputs back, weighted by the routing weights (``combine``). Experts only see a :class:`DispatchOutput`, so any
+dispatcher can be paired with any experts backend.
+"""
+
+from dataclasses import dataclass
+from typing import Any, List, Optional, Protocol
+
+import torch
+import torch.distributed as dist
+from torch import Tensor
+from torch.distributed import ProcessGroup
+
+from ._operation import all_to_all_uneven
+
+__all__ = ["DispatchOutput", "Dispatcher", "NCCLAllToAllDispatcher"]
+
+
+@dataclass
+class DispatchOutput:
+    """What a dispatcher gives the experts.
+
+    Attributes:
+        tokens: ``[S, H]`` rows for the local experts, grouped by local expert (all rows of local expert 0 first).
+            Within an expert, rows keep the order of their tokens in the batch of the whole EP group (source rank
+            first, then token position), so expert grads can be compared bitwise with a single-GPU reference.
+        tokens_per_expert: ``[E_local]`` int64 row count per local expert, on the same device as ``tokens``.
+        tokens_per_expert_cpu: the same counts as a list, when the dispatcher already has them on the host.
+        probs: ``[S]`` fp32 routing weight of each row, for dispatchers that let the experts apply the weights. ``None``
+            when ``combine`` applies them on the source rank. Experts never apply them otherwise.
+        handle: dispatcher-private state needed by ``combine``.
+    """
+
+    tokens: Tensor
+    tokens_per_expert: Tensor
+    tokens_per_expert_cpu: Optional[List[int]] = None
+    probs: Optional[Tensor] = None
+    handle: Any = None
+
+
+class Dispatcher(Protocol):
+    def dispatch(self, hidden_states: Tensor, topk_idx: Tensor, topk_weights: Tensor) -> DispatchOutput:
+        """Send ``hidden_states`` ``[T, H]`` to the ranks of the experts in ``topk_idx`` ``[T, K]`` (global ids)."""
+
+    def combine(self, expert_output: Tensor, dispatched: DispatchOutput) -> Tensor:
+        """Return ``[T, H]``: for every token, the sum of its expert outputs weighted by ``topk_weights``."""
+
+
+class NCCLAllToAllDispatcher:
+    """Dispatcher built on NCCL all-to-all, with one host sync per layer (the split sizes NCCL needs on the host).
+
+    The weighted sum is done on the source rank in fp32 with the fp32 routing weights and cast once, which matches
+    the grouped_mm path of HF transformers v5 bitwise in the forward pass.
+
+    Args:
+        num_experts: total number of experts in the layer.
+        ep_group: expert parallel group; experts are split evenly and contiguously over its ranks.
+    """
+
+    def __init__(self, num_experts: int, ep_group: Optional[ProcessGroup] = None):
+        self.ep_group = ep_group
+        self.ep_size = dist.get_world_size(ep_group) if dist.is_initialized() else 1
+        if num_experts % self.ep_size != 0:
+            raise ValueError(f"num_experts ({num_experts}) must be divisible by the ep size ({self.ep_size})")
+        self.num_experts = num_experts
+        self.num_local_experts = num_experts // self.ep_size
+
+    def dispatch(self, hidden_states: Tensor, topk_idx: Tensor, topk_weights: Tensor) -> DispatchOutput:
+        num_tokens, top_k = topk_idx.shape
+        flat_idx = topk_idx.reshape(-1)
+        # (token, slot) pairs sorted by global expert id, i.e. by destination rank then local expert
+        order = torch.sort(flat_idx, stable=True)[1]
+        send = hidden_states[order // top_k]
+        send_counts = torch.bincount(flat_idx, minlength=self.num_experts).view(self.ep_size, self.num_local_experts)
+        if self.ep_size > 1:
+            recv_counts = torch.empty_like(send_counts)
+            dist.all_to_all_single(recv_counts, send_counts, group=self.ep_group)  # [source rank, local expert]
+        else:
+            recv_counts = send_counts
+        # the one host sync: split sizes for NCCL and the per-expert row counts
+        sizes = torch.cat([send_counts.sum(1), recv_counts.sum(1), recv_counts.sum(0)]).tolist()
+        send_splits = sizes[: self.ep_size]
+        recv_splits = sizes[self.ep_size : 2 * self.ep_size]
+        tokens_per_expert_cpu = sizes[2 * self.ep_size :]
+        if self.ep_size > 1:
+            recv, _ = all_to_all_uneven(send, send_splits, recv_splits, self.ep_group)
+        else:
+            recv = send
+        # rows arrive grouped by (source rank, local expert); regroup them by local expert on the GPU
+        local_expert = torch.arange(self.num_local_experts, device=recv.device).repeat(self.ep_size)
+        row_expert = local_expert.repeat_interleave(recv_counts.reshape(-1), output_size=recv.size(0))
+        perm = torch.sort(row_expert, stable=True)[1]
+        return DispatchOutput(
+            tokens=recv[perm],
+            tokens_per_expert=recv_counts.sum(0),
+            tokens_per_expert_cpu=tokens_per_expert_cpu,
+            handle=(order, perm, send_splits, recv_splits, num_tokens, top_k, topk_weights),
+        )
+
+    def combine(self, expert_output: Tensor, dispatched: DispatchOutput) -> Tensor:
+        order, perm, send_splits, recv_splits, num_tokens, top_k, topk_weights = dispatched.handle
+        out = expert_output[_inverse_permutation(perm)]
+        if self.ep_size > 1:
+            out, _ = all_to_all_uneven(out, recv_splits, send_splits, self.ep_group)
+        weights = topk_weights.reshape(-1)[order].float()
+        weighted = (out.float() * weights[:, None])[_inverse_permutation(order)]
+        return weighted.view(num_tokens, top_k, out.size(-1)).sum(1).to(expert_output.dtype)
+
+
+def _inverse_permutation(perm: Tensor) -> Tensor:
+    inv = torch.empty_like(perm)
+    inv[perm] = torch.arange(perm.numel(), device=perm.device)
+    return inv

colossalai/moe/experts.py

@@ -0,0 +1,78 @@
+"""Experts backends: the SwiGLU MLPs of the local experts, run on the rows of a :class:`DispatchOutput`.
+
+Backends are stateless; the weights belong to the layer and use the transformers v5 layout:
+``gate_up_proj`` ``[E_local, 2 * I, H]`` (gate rows first) and ``down_proj`` ``[E_local, H, I]``.
+Outputs are not multiplied by the routing weights: the dispatcher's ``combine`` does that.
+"""
+
+from typing import Callable, Protocol
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+
+from .dispatcher import DispatchOutput
+
+__all__ = ["Experts", "GroupedMMExperts", "ReferenceExperts"]
+
+ActFn = Callable[[Tensor], Tensor]
+
+
+class Experts(Protocol):
+    def __call__(self, dispatched: DispatchOutput, gate_up_proj: Tensor, down_proj: Tensor, act_fn: ActFn) -> Tensor:
+        """Return ``[S, H]``, one output row per row of ``dispatched.tokens``, in the same order."""
+
+
+def _torch_version() -> tuple:
+    return tuple(int(v) for v in torch.__version__.split("+")[0].split(".")[:2])
+
+
+class GroupedMMExperts:
+    """All local experts in two ``torch._grouped_mm`` calls (torch >= 2.8).
+
+    Before torch 2.11, ``_grouped_mm`` needs every group's row count to be a multiple of 8 (16 bytes in bf16), so the
+    rows are padded per expert and the padding is dropped afterwards.
+    """
+
+    def __init__(self):
+        if _torch_version() < (2, 8) or not hasattr(torch, "_grouped_mm"):
+            raise RuntimeError(f"GroupedMMExperts needs torch >= 2.8, got {torch.__version__}; use ReferenceExperts")
+        self.align = 1 if _torch_version() >= (2, 11) else 8
+
+    def __call__(self, dispatched: DispatchOutput, gate_up_proj: Tensor, down_proj: Tensor, act_fn: ActFn) -> Tensor:
+        tokens, counts = dispatched.tokens, dispatched.tokens_per_expert
+        if self.align == 1:
+            return self._mlp(tokens, counts, gate_up_proj, down_proj, act_fn)
+        num_experts = counts.numel()
+        padded_counts = (counts + self.align - 1) // self.align * self.align
+        start = torch.cumsum(counts, 0) - counts
+        padded_start = torch.cumsum(padded_counts, 0) - padded_counts
+        row_expert = torch.arange(num_experts, device=tokens.device).repeat_interleave(
+            counts, output_size=tokens.size(0)
+        )
+        dest = torch.arange(tokens.size(0), device=tokens.device) - start[row_expert] + padded_start[row_expert]
+        padded = tokens.new_zeros(tokens.size(0) + (self.align - 1) * num_experts, tokens.size(1))
+        padded = padded.index_copy(0, dest, tokens)
+        return self._mlp(padded, padded_counts, gate_up_proj, down_proj, act_fn)[dest]
+
+    @staticmethod
+    def _mlp(tokens: Tensor, counts: Tensor, gate_up_proj: Tensor, down_proj: Tensor, act_fn: ActFn) -> Tensor:
+        offs = torch.cumsum(counts, 0).to(torch.int32)
+        gate, up = torch._grouped_mm(tokens, gate_up_proj.transpose(1, 2), offs=offs).chunk(2, dim=-1)
+        return torch._grouped_mm(act_fn(gate) * up, down_proj.transpose(1, 2), offs=offs)
+
+
+class ReferenceExperts:
+    """One ``F.linear`` pair per expert. Slow; for debugging and as the reference in tests."""
+
+    def __call__(self, dispatched: DispatchOutput, gate_up_proj: Tensor, down_proj: Tensor, act_fn: ActFn) -> Tensor:
+        counts = dispatched.tokens_per_expert_cpu
+        if counts is None:
+            counts = dispatched.tokens_per_expert.tolist()
+        outputs = []
+        for expert, rows in enumerate(dispatched.tokens.split(counts)):
+            gate, up = F.linear(rows, gate_up_proj[expert]).chunk(2, dim=-1)
+            outputs.append(F.linear(act_fn(gate) * up, down_proj[expert]))
+        if not outputs:
+            return dispatched.tokens.new_zeros(0, down_proj.size(1))
+        return torch.cat(outputs)