From c17a765572315739e4955f69b5d773e058a66b54 Mon Sep 17 00:00:00 2001 From: CaiLingchen <2589093246@qq.com> Date: Fri, 24 Jul 2026 16:43:47 +0800 Subject: [PATCH] =?UTF-8?q?EP=20experts=20=E6=94=B9=E7=94=A8=20EpExpertsGm?= =?UTF-8?q?m=20=E6=8E=A5=E5=8F=A3?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/twinkle/kernel/ops/__init__.py | 14 +++ src/twinkle/kernel/ops/ep/__init__.py | 112 ++++++++++++++++++ src/twinkle/kernel/ops/ep/loop.py | 66 +++++++++++ src/twinkle/kernel/ops/ep/npu.py | 103 ++++++++++++++++ .../model/transformers/moe/ep_utils.py | 22 +++- .../model/transformers/moe/expert_parallel.py | 49 +------- 6 files changed, 314 insertions(+), 52 deletions(-) create mode 100644 src/twinkle/kernel/ops/__init__.py create mode 100644 src/twinkle/kernel/ops/ep/__init__.py create mode 100644 src/twinkle/kernel/ops/ep/loop.py create mode 100644 src/twinkle/kernel/ops/ep/npu.py diff --git a/src/twinkle/kernel/ops/__init__.py b/src/twinkle/kernel/ops/__init__.py new file mode 100644 index 000000000..aa95951c3 --- /dev/null +++ b/src/twinkle/kernel/ops/__init__.py @@ -0,0 +1,14 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Backend-agnostic operator interfaces with per-platform implementations. + +Each op family defines an abstract base class (e.g. ``EpExpertsGmm``) plus a +dispatcher that picks the first eligible backend implementation from a +lazily-built registry. Callers invoke the dispatcher only; backend details +(NPU, GPU, ...) stay hidden behind the interface. +""" +from .ep import EpExpertsGmm, ep_forward + +__all__ = [ + 'EpExpertsGmm', + 'ep_forward', +] diff --git a/src/twinkle/kernel/ops/ep/__init__.py b/src/twinkle/kernel/ops/ep/__init__.py new file mode 100644 index 000000000..c011d3ced --- /dev/null +++ b/src/twinkle/kernel/ops/ep/__init__.py @@ -0,0 +1,112 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Grouped-matmul interface for expert-parallel (EP) experts compute. + +``EpExpertsGmm`` is the abstract backend interface; each platform provides a +subclass (e.g. ``NpuEpExpertsGmm`` in ``ep_experts_gmm_npu.py``) that is +registered lazily. The generic per-expert loop (``LoopEpExpertsGmm`` in +``ep_experts_gmm_loop.py``) is registered last as the always-eligible +fallback. ``ep_forward`` is the single entry point used by the EP +forward: it dispatches to the first eligible backend, so it always returns a +result. +""" +from __future__ import annotations + +import torch +from abc import ABC, abstractmethod +from torch import nn + +from twinkle import get_logger + +logger = get_logger() + + +class EpExpertsGmm(ABC): + """Backend interface for batched EP experts compute via grouped matmul. + + A backend computes all local experts in one shot from the permuted token + buffer (tokens sorted by expert, contiguous chunks per local expert), + replacing the per-expert Python loop. + """ + + name: str + # Fallback backends (e.g. the per-expert loop) are always eligible and are + # registered last; hitting one means every accelerated backend declined. + fallback: bool = False + + @abstractmethod + def ineligible_reason(self, experts_mod: nn.Module) -> str | None: + """Return why this backend cannot handle ``experts_mod``, or None if it can.""" + + @abstractmethod + def forward( + self, + experts_mod: nn.Module, + permuted_tokens: torch.Tensor, + num_global_sum_tokens_per_local_expert: torch.Tensor, + experts_per_rank: int, + ) -> torch.Tensor: + """Compute local experts on the permuted token buffer.""" + + +_IMPLS: list[EpExpertsGmm] | None = None +_PATH_LOGGED = False +_WARN_LOGGED = False + + +def _get_impls() -> list[EpExpertsGmm]: + """Build the backend registry on first use. + + Each backend module is imported defensively: platforms lacking its + dependencies (e.g. no torch_npu) simply skip that backend. + """ + global _IMPLS + if _IMPLS is None: + _IMPLS = [] + try: + from .npu import NpuEpExpertsGmm + _IMPLS.append(NpuEpExpertsGmm()) + except ImportError: + pass + # The loop fallback has no optional dependencies and is always last: + # it guarantees the dispatcher never runs out of backends. + from .loop import LoopEpExpertsGmm + _IMPLS.append(LoopEpExpertsGmm()) + return _IMPLS + + +def ep_forward( + experts_mod: nn.Module, + permuted_tokens: torch.Tensor, + num_global_sum_tokens_per_local_expert: torch.Tensor, + experts_per_rank: int, +) -> torch.Tensor: + """Dispatch to the first eligible backend (the loop fallback guarantees a hit). + + One-time logging: INFO on the first accelerated dispatch (process-wide); + WARNING with the ineligibility reasons when falling back to the loop + (once per experts instance, as reasons may differ per MoE block). + """ + if permuted_tokens.numel() == 0: + # Preserve the autograd edge to token_pre_all2all. Returning a new + # empty tensor can make this rank skip the matching backward + # all-to-all, causing EP collective order divergence. + return permuted_tokens + + global _PATH_LOGGED, _WARN_LOGGED + ineligible: list[str] = [] + for impl in _get_impls(): + reason = impl.ineligible_reason(experts_mod) + if reason is not None: + ineligible.append(f'{impl.name}: {reason}') + continue + if impl.fallback: + if ineligible and not _WARN_LOGGED: + detail = '; '.join(ineligible) + logger.warning(f'EP experts compute: grouped matmul disabled ({detail}); ' + f'falling back to {impl.name}.') + _WARN_LOGGED = True + elif not _PATH_LOGGED: + logger.info(f'EP experts compute: using {impl.name} grouped matmul (per-expert loop replaced).') + _PATH_LOGGED = True + return impl.forward(experts_mod, permuted_tokens, num_global_sum_tokens_per_local_expert, experts_per_rank) + raise RuntimeError('no EP experts backend registered (loop fallback missing)') diff --git a/src/twinkle/kernel/ops/ep/loop.py b/src/twinkle/kernel/ops/ep/loop.py new file mode 100644 index 000000000..c08d03040 --- /dev/null +++ b/src/twinkle/kernel/ops/ep/loop.py @@ -0,0 +1,66 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Generic per-expert Python loop backend for the EP experts interface. + +This is the always-eligible fallback: it runs on any platform (pure PyTorch) +and is registered last, so the dispatcher can never run out of backends. +""" +from __future__ import annotations + +import torch +import torch.nn.functional as F +from torch import nn + +from . import EpExpertsGmm + + +class LoopEpExpertsGmm(EpExpertsGmm): + """Per-expert Python loop over the permuted token buffer (generic fallback).""" + + name = 'per-expert loop' + fallback = True + + def ineligible_reason(self, experts_mod: nn.Module) -> str | None: + # The loop is the reference implementation: always eligible. + return None + + def forward( + self, + experts_mod: nn.Module, + permuted_tokens: torch.Tensor, + num_global_sum_tokens_per_local_expert: torch.Tensor, + experts_per_rank: int, + ) -> torch.Tensor: + input_dtype = permuted_tokens.dtype + + cumsum = torch.zeros(experts_per_rank + 1, dtype=torch.long) + for i in range(experts_per_rank): + cumsum[i + 1] = cumsum[i] + int(num_global_sum_tokens_per_local_expert[i].item()) + + output_chunks = [] + for i in range(experts_per_rank): + start = int(cumsum[i].item()) + end = int(cumsum[i + 1].item()) + expert_in = permuted_tokens[start:end] + if expert_in.numel() == 0: + output_chunks.append(expert_in) + continue + + gate_up = experts_mod.gate_up_proj[i] + down = experts_mod.down_proj[i] + compute_dtype = gate_up.dtype + if expert_in.dtype != compute_dtype: + expert_in = expert_in.to(compute_dtype) + gate_up_out = F.linear(expert_in, gate_up) + if hasattr(experts_mod, '_apply_gate'): + out = experts_mod._apply_gate(gate_up_out) + else: + gate, up = gate_up_out.chunk(2, dim=-1) + out = experts_mod.act_fn(gate) * up + out = F.linear(out, down) + + if out.dtype != input_dtype: + out = out.to(input_dtype) + output_chunks.append(out) + + return torch.cat( + output_chunks, dim=0) if output_chunks else permuted_tokens.new_empty(0, permuted_tokens.size(-1)) diff --git a/src/twinkle/kernel/ops/ep/npu.py b/src/twinkle/kernel/ops/ep/npu.py new file mode 100644 index 000000000..8fe215186 --- /dev/null +++ b/src/twinkle/kernel/ops/ep/npu.py @@ -0,0 +1,103 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""NPU backend for the EP experts grouped-matmul interface.""" +from __future__ import annotations + +import torch +from torch import nn + +from . import EpExpertsGmm + + +class NpuEpExpertsGmm(EpExpertsGmm): + """Batched EP experts compute via ``torch_npu.npu_grouped_matmul``.""" + + name = 'NPU' + + def ineligible_reason(self, experts_mod: nn.Module) -> str | None: + """Check whether the tensor-experts EP compute can use NPU grouped matmul. + + Falls back to the per-expert Python loop when: + - torch_npu / NPU is unavailable; + - there is no way to compute the gated activation: no ``_apply_gate`` + hook and a non-SiLU ``act_fn`` (``npu_swiglu`` == silu(gate) * up); + - the packed weights are not plain 3D floating tensors (e.g. DTensor + that was not unsharded, quantized weights, etc.). + + Note on ``_apply_gate``: transformers injects a default implementation + (``_default_apply_gate``: chunk + act_fn(gate) * up) onto many MoE + expert classes. With SiLU activation that is exactly ``npu_swiglu``, so + it does not block the fused path. A *custom* gate hook is still + supported — it is called eagerly on the grouped-matmul output. + """ + apply_gate = getattr(experts_mod, '_apply_gate', None) + is_default_gate = getattr(apply_gate, '__name__', None) == '_default_apply_gate' + if apply_gate is None or is_default_gate: + # Activation will be computed by npu_swiglu == silu(gate) * up. + act_fn = getattr(experts_mod, 'act_fn', None) + if act_fn is None: + return 'no act_fn attribute' + # nn.SiLU and transformers' SiLUActivation both compute F.silu + # exactly; either is equivalent to the silu inside npu_swiglu. + if act_fn.__class__.__name__ not in ('SiLU', 'SiLUActivation'): + return f'act_fn is {act_fn.__class__.__name__}, not SiLU' + # else: custom gate hook is called eagerly in the GMM path — allowed. + try: + import torch_npu # noqa: F401 + except ImportError: + return 'torch_npu not importable' + if not torch.npu.is_available(): + return 'NPU not available' + for name in ('gate_up_proj', 'down_proj'): + weight = getattr(experts_mod, name, None) + # Note: accept plain tensors and nn.Parameter. Under FSDP2 the + # pre-forward hook has already unsharded params by the time + # ep_forward runs; under PEFT target_parameters the parametrized + # property returns a merged plain tensor. Anything else (DTensor, + # quantized types) falls back to the loop. + if not isinstance(weight, torch.Tensor) or weight.__class__.__name__ == 'DTensor': + return f'{name} is {type(weight).__name__}, not a plain tensor' + if weight.ndim != 3 or not weight.dtype.is_floating_point: + return f'{name} has shape {tuple(weight.shape)} dtype {weight.dtype}' + return None + + def forward( + self, + experts_mod: nn.Module, + permuted_tokens: torch.Tensor, + num_global_sum_tokens_per_local_expert: torch.Tensor, + experts_per_rank: int, + ) -> torch.Tensor: + """Compute local experts with one grouped matmul instead of a Python loop. + + Mathematically equivalent to the per-expert loop: grouped matmul + applies each expert's weight to its contiguous token chunk, and + ``npu_swiglu(x) == silu(gate) * up`` for gate/up = x.chunk(2, -1). + The group list stays on device, eliminating the per-expert ``.item()`` + host synchronizations of the loop implementation. + """ + import torch_npu + + from twinkle.kernel.npu_impls.moe import GmmFunction, _get_cached_expert_weights + + input_dtype = permuted_tokens.dtype + hidden_dim = permuted_tokens.size(-1) + # counts arrives as a CPU tensor from preprocess(); grouped matmul needs + # the group list on device. + group_list = num_global_sum_tokens_per_local_expert.to(device=permuted_tokens.device, dtype=torch.int64) + gate_up_weight, down_weight = _get_cached_expert_weights(experts_mod, input_dtype, hidden_dim) + + expert_in = permuted_tokens + if expert_in.dtype != gate_up_weight.dtype: + expert_in = expert_in.to(gate_up_weight.dtype) + intermediate = GmmFunction.apply(expert_in, group_list, gate_up_weight) + apply_gate = getattr(experts_mod, '_apply_gate', None) + if apply_gate is not None and getattr(apply_gate, '__name__', '') != '_default_apply_gate': + # Custom gate hook: honor it eagerly instead of the fused swiglu. + activated = apply_gate(intermediate) + else: + # No hook, or transformers' default gate: silu(gate) * up == npu_swiglu. + activated = torch_npu.npu_swiglu(intermediate, dim=-1) + out = GmmFunction.apply(activated, group_list, down_weight) + if out.dtype != input_dtype: + out = out.to(input_dtype) + return out diff --git a/src/twinkle/model/transformers/moe/ep_utils.py b/src/twinkle/model/transformers/moe/ep_utils.py index f64c17964..bb2bed884 100644 --- a/src/twinkle/model/transformers/moe/ep_utils.py +++ b/src/twinkle/model/transformers/moe/ep_utils.py @@ -8,6 +8,8 @@ import torch.distributed as dist from typing import Optional +from twinkle import torch_util + # ========================== comm ========================== class _AllToAll(torch.autograd.Function): @@ -113,7 +115,11 @@ def permute(tokens: torch.Tensor, expert_mask: torch.Tensor): sorted_indices = token_indices.masked_select(expert_mask) # use the mapping to permute the tokens - permuted_input = tokens.index_select(0, sorted_indices) + # NOTE: use advanced indexing instead of index_select — index_select's + # backward (index_add) is broken on some torch_npu/CANN versions + # (aclnnIndexAdd fails for any dtype), while x[idx] backward (index_put) + # works. Mathematically identical. + permuted_input = tokens[sorted_indices] return permuted_input, sorted_indices @@ -203,6 +209,12 @@ def preprocess( ) dist.all_gather_into_tensor(num_global_tokens_per_expert, num_local_tokens_per_expert, group=ep_group) + # The collective may run on a separate stream (e.g. HCCL on NPU) whose + # completion is not always ordered with the host reads (.tolist()/.item()) + # below — observed in practice as garbage split sizes. Force a device sync + # before reading gathered results back on the host. + torch_util.synchronize() + # [ep_size, num_local_experts] start_idx, end_idx = rank * num_local_experts, (rank + 1) * num_local_experts num_global_tokens_per_local_expert = num_global_tokens_per_expert[:, start_idx:end_idx].contiguous() @@ -211,11 +223,13 @@ def preprocess( output_splits = num_global_tokens_per_local_expert.sum(dim=1).tolist() # [num_local_expert] - num_global_sum_tokens_per_local_expert = num_global_tokens_per_local_expert.sum(dim=0).to( - torch.device('cpu'), non_blocking=True) + # NOTE: keep these copies synchronous. With non_blocking=True the CPU + # tensors are read (tolist()/item()) before the async D2H copy lands, + # producing garbage split sizes (observed on NPU). + num_global_sum_tokens_per_local_expert = num_global_tokens_per_local_expert.sum(dim=0).to(torch.device('cpu')) num_global_tokens_per_local_expert = num_global_tokens_per_local_expert.view(-1, num_local_experts).to( - torch.device('cpu'), non_blocking=True) + torch.device('cpu')) return input_splits, output_splits, num_global_tokens_per_local_expert, num_global_sum_tokens_per_local_expert diff --git a/src/twinkle/model/transformers/moe/expert_parallel.py b/src/twinkle/model/transformers/moe/expert_parallel.py index f043e547c..218e7b337 100644 --- a/src/twinkle/model/transformers/moe/expert_parallel.py +++ b/src/twinkle/model/transformers/moe/expert_parallel.py @@ -4,11 +4,11 @@ import inspect import torch import torch.distributed as dist -import torch.nn.functional as F from dataclasses import dataclass from torch import nn from typing import Any, Dict, Iterable, List, Optional, Tuple +from twinkle.kernel.ops import ep_forward from twinkle.model.transformers.moe.ep_utils import preprocess, token_pre_all2all, tokens_post_all2all from twinkle.utils import DeviceMesh @@ -330,53 +330,6 @@ def _install_ep_forward(experts_mod: nn.Module, experts_per_rank: int) -> None: if getattr(experts_mod, '_ep_forward_installed', False): return - def ep_forward( - self, - permuted_tokens: torch.Tensor, - num_global_sum_tokens_per_local_expert: torch.Tensor, - experts_per_rank: int, - ) -> torch.Tensor: - if permuted_tokens.numel() == 0: - # Preserve the autograd edge to token_pre_all2all. Returning a new - # empty tensor can make this rank skip the matching backward - # all-to-all, causing EP collective order divergence. - return permuted_tokens - - input_dtype = permuted_tokens.dtype - - cumsum = torch.zeros(experts_per_rank + 1, dtype=torch.long) - for i in range(experts_per_rank): - cumsum[i + 1] = cumsum[i] + int(num_global_sum_tokens_per_local_expert[i].item()) - - output_chunks = [] - for i in range(experts_per_rank): - start = int(cumsum[i].item()) - end = int(cumsum[i + 1].item()) - expert_in = permuted_tokens[start:end] - if expert_in.numel() == 0: - output_chunks.append(expert_in) - continue - - gate_up = self.gate_up_proj[i] - down = self.down_proj[i] - compute_dtype = gate_up.dtype - if expert_in.dtype != compute_dtype: - expert_in = expert_in.to(compute_dtype) - gate_up_out = F.linear(expert_in, gate_up) - if hasattr(self, '_apply_gate'): - out = self._apply_gate(gate_up_out) - else: - gate, up = gate_up_out.chunk(2, dim=-1) - out = self.act_fn(gate) * up - out = F.linear(out, down) - - if out.dtype != input_dtype: - out = out.to(input_dtype) - output_chunks.append(out) - - return torch.cat( - output_chunks, dim=0) if output_chunks else permuted_tokens.new_empty(0, permuted_tokens.size(-1)) - import types experts_mod.forward = types.MethodType(ep_forward, experts_mod) experts_mod._ep_forward_installed = True