diff --git a/.github/scripts/ci/run_colossalai_batch.sh b/.github/scripts/ci/run_colossalai_batch.sh index f030d4ae76ef..e9c7c41e59ca 100644 --- a/.github/scripts/ci/run_colossalai_batch.sh +++ b/.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 diff --git a/.github/scripts/ci/run_cpu_unit_tests.sh b/.github/scripts/ci/run_cpu_unit_tests.sh index 777743d83702..1a124b45cfc7 100644 --- a/.github/scripts/ci/run_cpu_unit_tests.sh +++ b/.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 diff --git a/colossalai/moe/__init__.py b/colossalai/moe/__init__.py index f6ced3fedc5a..6b47b77e5915 100644 --- a/colossalai/moe/__init__.py +++ b/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}") diff --git a/colossalai/moe/dispatcher.py b/colossalai/moe/dispatcher.py new file mode 100644 index 000000000000..047ceb8ed93d --- /dev/null +++ b/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 diff --git a/colossalai/moe/experts.py b/colossalai/moe/experts.py new file mode 100644 index 000000000000..6f95d6151b45 --- /dev/null +++ b/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) diff --git a/colossalai/moe/layer.py b/colossalai/moe/layer.py new file mode 100644 index 000000000000..4e5d6745a0f9 --- /dev/null +++ b/colossalai/moe/layer.py @@ -0,0 +1,148 @@ +"""EPMoELayer: the expert-parallel part of an MoE block, shared by all models. + +An MoE block is split into Router -> dispatch -> experts -> combine. Models keep their own router (top-k selection, +grouped routing, ...) and hand its result to this layer, which runs the other three steps with a pluggable +:class:`~colossalai.moe.dispatcher.Dispatcher` and :class:`~colossalai.moe.experts.Experts` backend. +""" + +import math +from dataclasses import replace +from typing import Optional + +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor +from torch.distributed import ProcessGroup + +from colossalai.tensor.moe_tensor.api import ExpertGradScale, set_expert_grad_scale, set_moe_tensor_ep_group + +from ._operation import EPGradScalerIn, EPGradScalerOut +from .dispatcher import Dispatcher, NCCLAllToAllDispatcher +from .experts import ActFn, Experts, GroupedMMExperts + +__all__ = ["EPMoELayer"] + + +class EPMoELayer(nn.Module): + """Routed experts of an MoE block with expert parallelism. + + Each rank of ``ep_group`` holds ``num_experts / ep_size`` consecutive experts as two parameters in the + transformers v5 layout: ``gate_up_proj`` ``[E_local, 2 * intermediate_size, hidden_size]`` and ``down_proj`` + ``[E_local, hidden_size, intermediate_size]``. + + Gradient contract: the gradient of an expert weight on one rank sums the contributions of all tokens the EP group + sent to it. Where its ``1 / ep_size`` factor is applied is recorded on both parameters + (:class:`~colossalai.tensor.moe_tensor.api.ExpertGradScale`): + + - ``IN_FORWARD`` (default): this layer applies it in backward, like the existing EP models, so the current + plugins average expert grads correctly. + - ``IN_REDUCTION``: this layer does not; the gradient reduction must divide expert grads by + ``replica_group_size * ep_size``. Only use it with a plugin that reads this tag, or expert grads are ``ep_size`` + times too large. + + Args: + num_experts: total number of routed experts. + hidden_size: model hidden size. + intermediate_size: hidden size of each expert MLP. + ep_group: expert parallel group; ``None`` means the default group. + act_fn: activation on the gate projection (SwiGLU: ``act_fn(gate) * up``). + dispatcher: defaults to :class:`~colossalai.moe.dispatcher.NCCLAllToAllDispatcher`. + experts: defaults to :class:`~colossalai.moe.experts.GroupedMMExperts`. + expert_grad_scale: see the gradient contract above. + """ + + def __init__( + self, + num_experts: int, + hidden_size: int, + intermediate_size: int, + ep_group: Optional[ProcessGroup] = None, + act_fn: ActFn = F.silu, + dispatcher: Optional[Dispatcher] = None, + experts: Optional[Experts] = None, + expert_grad_scale: ExpertGradScale = ExpertGradScale.IN_FORWARD, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ): + super().__init__() + if ep_group is None and dist.is_initialized(): + ep_group = dist.group.WORLD + self.ep_group = ep_group + self.ep_size = dist.get_world_size(ep_group) if dist.is_initialized() else 1 + self.ep_rank = dist.get_rank(ep_group) if dist.is_initialized() else 0 + 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 + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.act_fn = act_fn + self.expert_grad_scale = ExpertGradScale(expert_grad_scale) + self.dispatcher = dispatcher if dispatcher is not None else NCCLAllToAllDispatcher(num_experts, ep_group) + self.experts = experts if experts is not None else GroupedMMExperts() + + factory = dict(device=device, dtype=dtype) + self.gate_up_proj = nn.Parameter( + torch.empty(self.num_local_experts, 2 * intermediate_size, hidden_size, **factory) + ) + self.down_proj = nn.Parameter(torch.empty(self.num_local_experts, hidden_size, intermediate_size, **factory)) + self.reset_parameters() + for param in (self.gate_up_proj, self.down_proj): + set_moe_tensor_ep_group(param, ep_group) + set_expert_grad_scale(param, self.expert_grad_scale) + + @property + def local_expert_range(self) -> range: + """Global ids of the experts held by this rank.""" + start = self.ep_rank * self.num_local_experts + return range(start, start + self.num_local_experts) + + def reset_parameters(self) -> None: + # each expert initialised like nn.Linear + for weight in (self.gate_up_proj, self.down_proj): + bound = 1 / math.sqrt(weight.size(-1)) + nn.init.uniform_(weight, -bound, bound) + + @torch.no_grad() + def load_full_weights(self, gate_up_proj: Tensor, down_proj: Tensor) -> None: + """Copy this rank's experts out of the weights of all experts (``[E, 2I, H]`` and ``[E, H, I]``). + + Only ``.data`` is written, so the parameter objects that optimizers and ZeRO already hold stay valid. + """ + local = slice(self.local_expert_range.start, self.local_expert_range.stop) + self.gate_up_proj.data.copy_(gate_up_proj[local]) + self.down_proj.data.copy_(down_proj[local]) + + def forward(self, hidden_states: Tensor, topk_idx: Tensor, topk_weights: Tensor) -> Tensor: + """Run the routed experts. + + Args: + hidden_states: ``[..., hidden_size]`` input of the MoE block. + topk_idx: ``[..., top_k]`` global expert ids chosen by the router for each token. + topk_weights: ``[..., top_k]`` routing weights, ideally fp32; gradients flow back through them. + + Returns: + Tensor of the shape and dtype of ``hidden_states``: per token, the weighted sum of its experts' outputs. + """ + shape = hidden_states.shape + top_k = topk_idx.size(-1) + dispatched = self.dispatcher.dispatch( + hidden_states.reshape(-1, self.hidden_size), topk_idx.reshape(-1, top_k), topk_weights.reshape(-1, top_k) + ) + scale_in_forward = self.expert_grad_scale == ExpertGradScale.IN_FORWARD and self.ep_size > 1 + if scale_in_forward: + # backward: divide the expert weight grads by ep_size, leave the input grads unchanged + dispatched = replace(dispatched, tokens=EPGradScalerIn.apply(dispatched.tokens, self.ep_size)) + output = self.experts(dispatched, self.gate_up_proj, self.down_proj, self.act_fn) + if scale_in_forward: + output = EPGradScalerOut.apply(output, self.ep_size) + return self.dispatcher.combine(output, dispatched).view(shape) + + def extra_repr(self) -> str: + return ( + f"num_experts={self.num_experts}, num_local_experts={self.num_local_experts}, " + f"hidden_size={self.hidden_size}, intermediate_size={self.intermediate_size}, ep_size={self.ep_size}, " + f"expert_grad_scale={self.expert_grad_scale.value}" + ) diff --git a/colossalai/tensor/moe_tensor/api.py b/colossalai/tensor/moe_tensor/api.py index a565ecf34eb1..dc010fb5cd3c 100644 --- a/colossalai/tensor/moe_tensor/api.py +++ b/colossalai/tensor/moe_tensor/api.py @@ -1,4 +1,5 @@ -from typing import List +from enum import Enum +from typing import List, Optional import torch import torch.distributed as dist @@ -57,6 +58,46 @@ def is_moe_tensor_ep_fused(tensor: torch.Tensor) -> bool: return getattr(tensor, "ep_fused", False) +class ExpertGradScale(str, Enum): + """Where the ``1 / ep_size`` factor of an expert parameter's gradient is applied. + + The gradient of an expert weight on one rank sums the contributions of every token the EP group routed to it, so + averaging it over the data-parallel replicas needs one more division by ``ep_size``. + """ + + # inside the layer's autograd (EPGradScaler); gradient reduction divides by the replica group size only + IN_FORWARD = "in_forward" + # not in the layer; gradient reduction divides by replica group size * ep_size + IN_REDUCTION = "in_reduction" + + +def set_expert_grad_scale(tensor: torch.Tensor, scale: ExpertGradScale) -> None: + """Record where the ``1 / ep_size`` factor of this expert parameter's gradient is applied. + + Args: + tensor (torch.Tensor): An expert parameter, already marked with :func:`set_moe_tensor_ep_group`. + scale (ExpertGradScale): Where the factor is applied. + """ + assert is_moe_tensor(tensor), "call set_moe_tensor_ep_group first" + tensor.__setattr__("expert_grad_scale", ExpertGradScale(scale)) + + +def get_expert_grad_scale(tensor: torch.Tensor) -> Optional[ExpertGradScale]: + """ + Get where the ``1 / ep_size`` factor of the given parameter's gradient is applied. + + Args: + tensor (torch.Tensor): The tensor to be checked. + + Returns: + Optional[ExpertGradScale]: ``None`` for non-expert tensors. Expert tensors without a tag are treated as + ``IN_FORWARD``, which is what the existing EP models do. + """ + if not is_moe_tensor(tensor): + return None + return getattr(tensor, "expert_grad_scale", ExpertGradScale.IN_FORWARD) + + def get_moe_info(ep_size: int, dp_size: int, pp_size: int, ep_inside: bool) -> MoeParallelInfo: """ Get moe info for the given tensor. diff --git a/tests/test_moe/test_ep_moe_layer.py b/tests/test_moe/test_ep_moe_layer.py new file mode 100644 index 000000000000..a1c35b1fcf07 --- /dev/null +++ b/tests/test_moe/test_ep_moe_layer.py @@ -0,0 +1,204 @@ +import pytest +import torch +import torch.distributed as dist +import torch.nn.functional as F + +import colossalai +from colossalai.moe import EPMoELayer, GroupedMMExperts, NCCLAllToAllDispatcher, ReferenceExperts +from colossalai.tensor.moe_tensor.api import ( + ExpertGradScale, + get_ep_group, + get_ep_size, + get_expert_grad_scale, + is_moe_tensor, +) +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn + +# Relative L2 error allowed between the bf16 layer and an fp32 single-GPU reference. The largest error over all checks +# here is 6e-3; dropping the EP grad scaler, for example, gives 0.9. +RTOL = 2e-2 + +NUM_EXPERTS, TOP_K, HIDDEN_SIZE, INTERMEDIATE_SIZE, NUM_TOKENS = 8, 2, 64, 32, 64 + + +def ep_groups(ep_size: int): + """Split the world into consecutive groups of ep_size ranks; return this rank's group.""" + world, rank = dist.get_world_size(), dist.get_rank() + mine = None + for start in range(0, world, ep_size): + group = dist.new_group(list(range(start, start + ep_size))) + if start <= rank < start + ep_size: + mine = group + return mine + + +def route(x: torch.Tensor, gate: torch.Tensor, bias: torch.Tensor): + """Softmax top-k router with renormalised fp32 weights (Mixtral / Qwen style).""" + probs = (F.linear(x.float(), gate) + bias).softmax(-1) + weights, idx = torch.topk(probs, TOP_K, dim=-1) + return idx, weights / weights.sum(-1, keepdim=True) + + +def rel(a: torch.Tensor, b: torch.Tensor) -> float: + return ((a.float() - b.float()).norm() / b.float().norm().clamp_min(1e-12)).item() + + +@parameterize("ep_size", [1, 2, 4]) +def run_dispatch_contract(ep_size: int): + """Rows reach the right experts, grouped by expert and in global token order; combine undoes dispatch.""" + group = ep_groups(ep_size) + ep_rank = dist.get_rank(group) + dispatcher = NCCLAllToAllDispatcher(NUM_EXPERTS, group) + num_local = NUM_EXPERTS // ep_size + + # ep rank 1 sends nothing, and expert 1 is never chosen + num_tokens = 0 if ep_rank == 1 else NUM_TOKENS + generator = torch.Generator(device="cuda").manual_seed(dist.get_rank()) + choices = torch.tensor([e for e in range(NUM_EXPERTS) if e != 1], device="cuda") + scores = torch.rand(num_tokens, choices.numel(), device="cuda", generator=generator) + topk_idx = choices[scores.argsort(-1)[:, :TOP_K]] + topk_weights = torch.rand(num_tokens, TOP_K, device="cuda", generator=generator) + topk_weights = topk_weights / topk_weights.sum(-1, keepdim=True) + x = torch.randn(num_tokens, HIDDEN_SIZE, device="cuda", generator=generator) + x[:, 0] = torch.arange(num_tokens, device="cuda") + ep_rank * NUM_TOKENS # global token id within the ep group + + dispatched = dispatcher.dispatch(x, topk_idx, topk_weights) + counts = dispatched.tokens_per_expert.tolist() + assert dispatched.tokens_per_expert_cpu == counts + assert dispatched.tokens.size(0) == sum(counts) + + # expected rows of each local expert: ids of the tokens that chose it, in ep-group token order + gathered = [None] * ep_size + dist.all_gather_object(gathered, (x[:, 0].long().tolist(), topk_idx.tolist()), group=group) + for local, rows in enumerate(dispatched.tokens.split(counts)): + expert = ep_rank * num_local + local + expected = [tid for ids, idx in gathered for tid, chosen in zip(ids, idx) if expert in chosen] + assert rows[:, 0].long().tolist() == expected, f"expert {expert}: rows out of place" + + # identity experts: combine must give back sum_j w_j * x = x + y = dispatcher.combine(dispatched.tokens, dispatched) + assert y.shape == x.shape + torch.testing.assert_close(y, x, rtol=1e-6, atol=1e-6) + + +@parameterize("ep_size", [1, 2, 4]) +@parameterize("backend", ["grouped_mm", "reference"]) +@parameterize("grad_scale", [ExpertGradScale.IN_FORWARD, ExpertGradScale.IN_REDUCTION]) +@parameterize("skewed", [False, True]) +def run_layer_vs_reference(ep_size: int, backend: str, grad_scale: ExpertGradScale, skewed: bool): + """Output, input grad and expert weight grads against an fp32 single-GPU computation.""" + group = ep_groups(ep_size) + experts = GroupedMMExperts() if backend == "grouped_mm" else ReferenceExperts() + layer = EPMoELayer( + NUM_EXPERTS, + HIDDEN_SIZE, + INTERMEDIATE_SIZE, + ep_group=group, + experts=experts, + expert_grad_scale=grad_scale, + device="cuda", + dtype=torch.bfloat16, + ) + + generator = torch.Generator(device="cuda").manual_seed(0) # same weights on every rank + gate = torch.randn(NUM_EXPERTS, HIDDEN_SIZE, device="cuda", generator=generator) * HIDDEN_SIZE**-0.5 + full_gate_up = torch.randn(NUM_EXPERTS, 2 * INTERMEDIATE_SIZE, HIDDEN_SIZE, device="cuda", generator=generator) * ( + HIDDEN_SIZE**-0.5 + ) + full_down = torch.randn(NUM_EXPERTS, HIDDEN_SIZE, INTERMEDIATE_SIZE, device="cuda", generator=generator) * ( + INTERMEDIATE_SIZE**-0.5 + ) + # skewed: the second half of the experts receives no token at all + bias = torch.zeros(NUM_EXPERTS, device="cuda") + if skewed: + bias[NUM_EXPERTS // 2 :] = -1e4 + layer.load_full_weights(full_gate_up, full_down) + + generator = torch.Generator(device="cuda").manual_seed(1 + dist.get_rank()) + x = torch.randn(NUM_TOKENS, HIDDEN_SIZE, device="cuda", generator=generator) + upstream = torch.randn(NUM_TOKENS, HIDDEN_SIZE, device="cuda", generator=generator) + + # layer (bf16), router in fp32 as models do + x_layer = x.to(torch.bfloat16).requires_grad_() + topk_idx, topk_weights = route(x_layer, gate, bias) + y = layer(x_layer, topk_idx, topk_weights) + (y.float() * upstream).sum().backward() + + # fp32 reference over this rank's tokens; expert grads summed over the ep group, as the layer computes them + x_ref = x.to(torch.bfloat16).float().requires_grad_() + gate_up_ref = full_gate_up.to(torch.bfloat16).float().requires_grad_() + down_ref = full_down.to(torch.bfloat16).float().requires_grad_() + ref_idx, ref_weights = route(x_ref, gate, bias) + y_ref = torch.zeros_like(x_ref) + for expert in range(NUM_EXPERTS): + token, slot = torch.where(ref_idx == expert) + g, u = F.linear(x_ref[token], gate_up_ref[expert]).chunk(2, dim=-1) + out = F.linear(F.silu(g) * u, down_ref[expert]) * ref_weights[token, slot, None] + y_ref = y_ref.index_add(0, token, out) + (y_ref * upstream).sum().backward() + dist.all_reduce(gate_up_ref.grad, group=group) + dist.all_reduce(down_ref.grad, group=group) + local = slice(layer.local_expert_range.start, layer.local_expert_range.stop) + + # IN_FORWARD divides expert grads by ep_size inside the layer; IN_REDUCTION leaves that to the reduction + factor = ep_size if grad_scale == ExpertGradScale.IN_FORWARD else 1 + config = f"ep={ep_size} {backend} {grad_scale.value} skewed={skewed}" + assert rel(y, y_ref) < RTOL, f"{config}: output {rel(y, y_ref):.3e}" + assert rel(x_layer.grad, x_ref.grad) < RTOL, f"{config}: input grad {rel(x_layer.grad, x_ref.grad):.3e}" + for name, grad, ref in [ + ("gate_up_proj", layer.gate_up_proj.grad, gate_up_ref.grad[local]), + ("down_proj", layer.down_proj.grad, down_ref.grad[local]), + ]: + assert grad is not None, f"{config}: {name} has no grad" + for expert, (g, r) in enumerate(zip(grad, ref)): + if r.abs().max() == 0: # expert received no token anywhere in the ep group + assert g.abs().max() == 0, f"{config}: {name}[{expert}] should be zero" + else: + err = rel(g.float() * factor, r) + assert err < RTOL, f"{config}: {name}[{expert}] grad {err:.3e}" + + +@parameterize("use_default_group", [False, True]) +def run_param_contract(use_default_group: bool): + """Expert params carry their EP tags, and loading weights keeps the parameter objects.""" + group = None if use_default_group else ep_groups(2) + expected_group = dist.group.WORLD if use_default_group else group + ep_size = dist.get_world_size(expected_group) + ep_rank = dist.get_rank(expected_group) + layer = EPMoELayer(NUM_EXPERTS, HIDDEN_SIZE, INTERMEDIATE_SIZE, ep_group=group, device="cuda") + assert layer.ep_group is expected_group + assert layer.dispatcher.ep_group is expected_group + assert layer.num_local_experts == NUM_EXPERTS // ep_size + start = ep_rank * layer.num_local_experts + assert layer.local_expert_range == range(start, start + layer.num_local_experts) + params = (layer.gate_up_proj, layer.down_proj) + for param in params: + assert is_moe_tensor(param) and get_ep_group(param) is expected_group + assert get_ep_size(param) == ep_size + assert get_expert_grad_scale(param) == ExpertGradScale.IN_FORWARD + ids = [id(p) for p in params] + layer.load_full_weights( + torch.ones(NUM_EXPERTS, 2 * INTERMEDIATE_SIZE, HIDDEN_SIZE, device="cuda"), + torch.ones(NUM_EXPERTS, HIDDEN_SIZE, INTERMEDIATE_SIZE, device="cuda"), + ) + assert [id(p) for p in (layer.gate_up_proj, layer.down_proj)] == ids + assert get_expert_grad_scale(torch.zeros(1)) is None + + +def run_dist(rank, world_size, port): + colossalai.launch(rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl") + torch.cuda.set_device(rank) + run_dispatch_contract() + run_layer_vs_reference() + run_param_contract() + + +@pytest.mark.dist +@pytest.mark.parametrize("world_size", [4]) +@rerun_if_address_is_in_use() +def test_ep_moe_layer(world_size): + spawn(run_dist, world_size) + + +if __name__ == "__main__": + test_ep_moe_layer(world_size=4) diff --git a/tests/test_moe/test_ep_moe_layer_cpu.py b/tests/test_moe/test_ep_moe_layer_cpu.py new file mode 100644 index 000000000000..3a716d4150a3 --- /dev/null +++ b/tests/test_moe/test_ep_moe_layer_cpu.py @@ -0,0 +1,29 @@ +import pytest +import torch.distributed as dist + +from colossalai.moe import EPMoELayer, ReferenceExperts +from colossalai.tensor.moe_tensor.api import ExpertGradScale, get_ep_group, get_ep_size, get_expert_grad_scale + + +@pytest.mark.parametrize("initialized", [False, True]) +def test_ep_moe_layer_default_process_group(initialized: bool, tmp_path): + """Default-group expert tags work with the parameter API without changing local construction.""" + assert not dist.is_initialized() + if initialized: + dist.init_process_group("gloo", init_method=(tmp_path / "rendezvous").as_uri(), rank=0, world_size=1) + try: + layer = EPMoELayer(4, 8, 4, experts=ReferenceExperts(), device="cpu") + expected_group = dist.group.WORLD if initialized else None + assert layer.ep_group is expected_group + assert layer.dispatcher.ep_group is expected_group + assert layer.ep_size == 1 + assert layer.ep_rank == 0 + assert layer.num_local_experts == 4 + for param in layer.parameters(): + assert get_ep_group(param) is expected_group + assert get_expert_grad_scale(param) == ExpertGradScale.IN_FORWARD + if initialized: + assert get_ep_size(param) == 1 + finally: + if initialized: + dist.destroy_process_group()