[Feat] Enable expert parallel for diffusion MoE layers (#1323)

Signed-off-by: Semmer2 <semmer@live.cn>

vllm_omni/diffusion/distributed/parallel_state.py

@@ -35,6 +35,8 @@ from vllm.distributed.parallel_state import get_tensor_model_parallel_world_size
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion import envs
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.forward_context import get_forward_context
 from vllm_omni.platforms import current_omni_platform
 
 from .group_coordinator import (
@@ -51,7 +53,7 @@ logger = init_logger(__name__)
 
 
 _WORLD: GroupCoordinator | None = None
-# get _TP from vllm.distributed.parallel_state
+# get _TP&_EP from vllm.distributed.parallel_state
 _SP: SequenceParallelGroupCoordinator | None = None
 _PP: PipelineGroupCoordinator | None = None
 _CFG: GroupCoordinator | None = None
@@ -188,6 +190,7 @@ class RankGenerator:
         self.fs = fs
         self.rank_offset = rank_offset
         self.world_size = tp * sp * pp * cfg * dp
+        self.ep = tp * sp * cfg * dp  # EP level exclude PP
 
         self.name_to_size = {
             "tp": self.tp,
@@ -250,6 +253,15 @@ class RankGenerator:
                 ranks.append(group)
             return ranks
 
+        if token == "ep":
+            ranks = []
+            num_pp_stages = self.pp
+            for i in range(num_pp_stages):
+                start = i * self.ep + self.rank_offset
+                end = start + self.ep
+                ranks.append(list(range(start, end)))
+            return ranks
+
         mask = self.get_mask(self.order, token)
         ranks = generate_masked_orthogonal_rank_groups(self.world_size, self.ordered_size, mask)
         if self.rank_offset > 0:
@@ -474,6 +486,7 @@ def init_model_parallel_group(
         "data",
         "pipeline",
         "tensor",
+        "expert",
         "sequence",
         "classifier_free_guidance",
         "fully_shard",
@@ -670,6 +683,7 @@ def initialize_model_parallel(
     pipeline_parallel_size: int = 1,
     fully_shard_degree: int = 1,
     hsdp_replicate_size: int = 1,
+    enable_expert_parallel: bool = False,
     backend: str | None = None,
 ) -> None:
     if backend is None:
@@ -837,6 +851,18 @@ def initialize_model_parallel(
         parallel_mode="fully_shard",
     )
 
+    if enable_expert_parallel:
+        od_config: OmniDiffusionConfig | None = get_forward_context().omni_diffusion_config
+        if od_config and od_config.is_moe:
+            vllm_parallel_state._EP = init_model_parallel_group(
+                group_ranks=rank_generator.get_ranks("ep"),
+                local_rank=get_world_group().local_rank,
+                backend=backend,
+                parallel_mode="expert",
+            )
+        else:
+            raise RuntimeError("Expert parallelism enabled for a non-MoE model ")
+
     init_dit_group(dit_parallel_size, backend)
 
 
@@ -861,6 +887,10 @@ def destroy_model_parallel():
         vllm_parallel_state._TP.destroy()
     vllm_parallel_state._TP = None
 
+    if vllm_parallel_state._EP:
+        vllm_parallel_state._EP.destroy()
+    vllm_parallel_state._EP = None
+
     global _PP
     if _PP:
         _PP.destroy()