From 6ac85d4607b0621703050844d5ba69b0adb0db59 Mon Sep 17 00:00:00 2001 From: Ting FU Date: Wed, 11 Mar 2026 10:45:36 +0800 Subject: [PATCH] [Feat] Enable expert parallel for diffusion MoE layers (#1323) Signed-off-by: Semmer2 --- .../diffusion/parallelism_acceleration.md | 54 +++- .../image_to_image/image_edit.py | 8 +- .../text_to_image/text_to_image.py | 8 +- .../text_to_video/text_to_video.py | 8 +- .../offline_inference/test_expert_parallel.py | 248 ++++++++++++++++++ vllm_omni/diffusion/data.py | 14 + .../diffusion/distributed/parallel_state.py | 32 ++- .../diffusion/worker/diffusion_worker.py | 2 + vllm_omni/entrypoints/async_omni.py | 2 + vllm_omni/entrypoints/async_omni_diffusion.py | 1 + vllm_omni/entrypoints/omni_diffusion.py | 2 +- 11 files changed, 361 insertions(+), 18 deletions(-) create mode 100644 tests/e2e/offline_inference/test_expert_parallel.py diff --git a/docs/user_guide/diffusion/parallelism_acceleration.md b/docs/user_guide/diffusion/parallelism_acceleration.md index 04edc59b3..bd2a58799 100644 --- a/docs/user_guide/diffusion/parallelism_acceleration.md +++ b/docs/user_guide/diffusion/parallelism_acceleration.md @@ -18,23 +18,26 @@ The following parallelism methods are currently supported in vLLM-Omni: 6. [HSDP](#hsdp): Hybrid Sharded Data Parallel shards model weights across GPUs using PyTorch FSDP2. This reduces per-GPU memory usage, enabling inference of large models on GPUs with limited memory. +7. [Expert Parallel](#expert-parallelism): Expert Parallelism shards the Experts of a Mixture-of-Experts (MoE) layer across multiple devices. During the forward, a gating mechanism routes tokens to their designated experts, necessitating cross-cards communication(all-to-all) to dispatch tokens to the correct ranks and combine the results. This parallelism allows for massive scaling of model parameters without a proportional increase in the computational load per device. + The following table shows which models are currently supported by parallelism method: ### ImageGen -| Model | Model Identifier | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel | VAE-Patch-Parallel | -|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:|:------------------:| -| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ✅ | ✅ | ❌ | ✅ | ❌ | -| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ✅ | ✅ | ❌ | ✅ | ❌ | -| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ | ❌ | ❌ | -| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ | -| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ❌ | -| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ❌ | -| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ✅ | ❌ | -| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ✅ (TP=2 only) | ✅ | -| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ✅ | ✅ | -| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ❌ | ❌ | ✅ | ❌ | -| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ❌ | ✅ | ✅ | ❌ | +| Model | Model Identifier | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel | VAE-Patch-Parallel | Expert-Parallel | +|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:|:------------------:|:---------------:| +| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ✅ | ✅ | ❌ | ✅ | ❌ | N/A | +| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ✅ | ✅ | ❌ | ✅ | ❌ | N/A | +| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ | ❌ | ❌ | N/A | +| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ | N/A | +| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A | +| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A | +| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A | +| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ✅ (TP=2 only) | ✅ | N/A | +| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ✅ | ✅ | N/A | +| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ❌ | ❌ | ✅ | ❌ | N/A | +| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ❌ | ✅ | ✅ | ❌ | N/A | +| **HunyuanImage3.0** | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | !!! note "TP Limitations for Diffusion Models" We currently implement Tensor Parallelism (TP) only for the DiT (Diffusion Transformer) blocks. This is because the `text_encoder` component in vLLM-Omni uses the original Transformers implementation, which does not yet support TP. @@ -414,3 +417,28 @@ vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091 --use-hsdp --usp #### Adding HSDP Support to New Models For detailed instructions on adding HSDP support to new models, see the [HSDP Contributing Guide](../../design/feature/hsdp.md). + +### Expert Parallelism + +Unlike Tensor Parallelism which shards every layer's weights, EP only shards the MoE expert MLP blocks. This significantly reduces the memory footprint of MoE models (e.g., HunyuanImage3.0) while maintaining constant dense-equivalent compute efficiency. Expert Parallelism is enabled via `DiffusionParallelConfig.enable_expert_parallel`. And `self.ep = tp * sp * cfg * dp` for now, so at least one of TP/SP/CFG/DP should set when EP enabled. + +#### Offline Inference + +```python +from vllm_omni import Omni +from vllm_omni.diffusion.data import DiffusionParallelConfig + +omni = Omni( + model="tencent/HunyuanImage-3.0", + parallel_config=DiffusionParallelConfig(tensor_parallel_size=8, enable_expert_parallel=True), +) + +outputs = omni.generate( + "A brown and white dog is running on the grass", + OmniDiffusionSamplingParams( + num_inference_steps=50, + width=1024, + height=1024, + ), +) +``` diff --git a/examples/offline_inference/image_to_image/image_edit.py b/examples/offline_inference/image_to_image/image_edit.py index 81fe24e21..83ac5af05 100644 --- a/examples/offline_inference/image_to_image/image_edit.py +++ b/examples/offline_inference/image_to_image/image_edit.py @@ -208,6 +208,11 @@ def parse_args() -> argparse.Namespace: default=1, help="Number of GPUs used for tensor parallelism (TP) inside the DiT.", ) + parser.add_argument( + "--enable-expert-parallel", + action="store_true", + help="Enable expert parallelism for MoE layers.", + ) parser.add_argument("--layers", type=int, default=4, help="Number of layers to decompose the input image into.") parser.add_argument( "--resolution", @@ -348,6 +353,7 @@ def main(): ring_degree=args.ring_degree, cfg_parallel_size=args.cfg_parallel_size, tensor_parallel_size=args.tensor_parallel_size, + enable_expert_parallel=args.enable_expert_parallel, ) # Configure cache based on backend type @@ -402,7 +408,7 @@ def main(): else: print(f" Input image size: {input_image.size}") print( - f" Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size}" + f" Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size}, enable_expert_parallel: {args.enable_expert_parallel}" ) print(f"{'=' * 60}\n") diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index 382971606..6f98403b9 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -170,6 +170,11 @@ def parse_args() -> argparse.Namespace: default=1, help="Number of GPUs used for tensor parallelism (TP) inside the DiT.", ) + parser.add_argument( + "--enable-expert-parallel", + action="store_true", + help="Enable expert parallelism for MoE layers.", + ) parser.add_argument( "--lora-path", type=str, @@ -258,6 +263,7 @@ def main(): cfg_parallel_size=args.cfg_parallel_size, tensor_parallel_size=args.tensor_parallel_size, vae_patch_parallel_size=args.vae_patch_parallel_size, + enable_expert_parallel=args.enable_expert_parallel, ) # Check if profiling is requested via environment variable @@ -323,7 +329,7 @@ def main(): print( f" Parallel configuration: tensor_parallel_size={args.tensor_parallel_size}, " f"ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, " - f"vae_patch_parallel_size={args.vae_patch_parallel_size}" + f"vae_patch_parallel_size={args.vae_patch_parallel_size}, enable_expert_parallel={args.enable_expert_parallel}." ) print(f" CPU offload: {args.enable_cpu_offload}") print(f" Image size: {args.width}x{args.height}") diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index c1efd342e..d7274e22b 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -127,6 +127,11 @@ def parse_args() -> argparse.Namespace: default=1, help="Number of GPUs used for VAE patch/tile parallelism (decode).", ) + parser.add_argument( + "--enable-expert-parallel", + action="store_true", + help="Enable expert parallelism for MoE layers.", + ) return parser.parse_args() @@ -160,6 +165,7 @@ def main(): cfg_parallel_size=args.cfg_parallel_size, tensor_parallel_size=args.tensor_parallel_size, vae_patch_parallel_size=args.vae_patch_parallel_size, + enable_expert_parallel=args.enable_expert_parallel, ) # Check if profiling is requested via environment variable @@ -193,7 +199,7 @@ def main(): print( f" Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}," f" cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size}," - f" vae_patch_parallel_size={args.vae_patch_parallel_size}" + f" vae_patch_parallel_size={args.vae_patch_parallel_size}, enable_expert_parallel={args.enable_expert_parallel}" ) print(f" Video size: {args.width}x{args.height}") print(f"{'=' * 60}\n") diff --git a/tests/e2e/offline_inference/test_expert_parallel.py b/tests/e2e/offline_inference/test_expert_parallel.py new file mode 100644 index 000000000..61f93e78f --- /dev/null +++ b/tests/e2e/offline_inference/test_expert_parallel.py @@ -0,0 +1,248 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +System test for Expert Parallel (EP). + +Tests verify that EP inference produces correct outputs compared to baseline. +""" + +import gc +import os +import time +from typing import NamedTuple + +import numpy as np +import pytest +import torch +import torch.distributed as dist +from PIL import Image + +from tests.utils import hardware_test +from vllm_omni import Omni +from vllm_omni.diffusion.data import DiffusionParallelConfig +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.platforms import current_omni_platform + +# Test configuration +MODELS = ["tencent/HunyuanImage-3.0"] +PROMPT = "A brown and white dog is running on the grass" +DEFAULT_HEIGHT = 1024 +DEFAULT_WIDTH = 1024 +DEFAULT_SEED = 1234 +DEFAULT_STEPS = 50 +DEFAULT_GUIDANCE_SCALE = 0.0 +COS_SIM_MEAN_THRESHOLD = 1e-2 +COS_SIM_MAX_THRESHOLD = 1e-2 +MSE_THRESHOLD = 1e-2 +DIFF_MEAN_THRESHOLD = 5e-2 +DIFF_MAX_THRESHOLD = 1 +TOTAL_CARD_NUM = 8 + + +class InferenceResult(NamedTuple): + """Result of an inference run.""" + + images: list[Image.Image] + elapsed_ms: float + + +def _cleanup_distributed(): + """Clean up distributed environment and GPU resources.""" + if dist.is_initialized(): + dist.destroy_process_group() + + for key in ["MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"]: + os.environ.pop(key, None) + + gc.collect() + if current_omni_platform.is_available(): + current_omni_platform.empty_cache() + current_omni_platform.synchronize() + + time.sleep(5) + + +def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float, float, float, float]: + """Return (mean_abs_diff, max_abs_diff) over RGB pixels in [0, 1].""" + ta = torch.from_numpy(np.asarray(a.convert("RGB"), dtype=np.float32) / 255.0) + tb = torch.from_numpy(np.asarray(b.convert("RGB"), dtype=np.float32) / 255.0) + assert ta.shape == tb.shape, f"Image shapes differ: {ta.shape} vs {tb.shape}" + abs_diff = torch.abs(ta - tb) + + vec_a = ta.reshape(-1) + vec_b = tb.reshape(-1) + cos_sim = torch.nn.functional.cosine_similarity(vec_a, vec_b, dim=0, eps=1e-8) + mse = torch.nn.functional.mse_loss(ta, tb) + return abs_diff.mean().item(), abs_diff.max().item(), cos_sim.mean().item(), cos_sim.max().item(), mse.item() + + +def _run_inference( + model_name: str, + tensor_parallel_size: int, + enable_expert_parallel: bool, + guidance_scale: float = DEFAULT_GUIDANCE_SCALE, + height: int = DEFAULT_HEIGHT, + width: int = DEFAULT_WIDTH, + seed: int = DEFAULT_SEED, + warmup: bool = True, +) -> InferenceResult: + """Run inference with specified configuration. + + Args: + warmup: If True, run one warmup iteration before the timed run. + """ + parallel_config = DiffusionParallelConfig( + tensor_parallel_size=tensor_parallel_size, + enable_expert_parallel=enable_expert_parallel, + ) + omni = Omni(model=model_name, parallel_config=parallel_config) + + try: + # Warmup run (not timed) + if warmup: + _ = omni.generate( + PROMPT, + OmniDiffusionSamplingParams( + height=height, + width=width, + num_inference_steps=DEFAULT_STEPS, + guidance_scale=guidance_scale, + generator=torch.Generator(current_omni_platform.device_type).manual_seed(seed), + num_outputs_per_prompt=1, + ), + ) + + # Timed run + start = time.time() + outputs = omni.generate( + PROMPT, + OmniDiffusionSamplingParams( + height=height, + width=width, + num_inference_steps=DEFAULT_STEPS, + guidance_scale=guidance_scale, + generator=torch.Generator(current_omni_platform.device_type).manual_seed(seed), + num_outputs_per_prompt=1, + ), + ) + elapsed_ms = (time.time() - start) * 1000 + + return InferenceResult( + images=outputs[0].request_output[0].images, + elapsed_ms=elapsed_ms, + ) + finally: + omni.close() + _cleanup_distributed() + + +# EP configurations: (tensor_parallel_size, enable_ep, height, width, is_perf_test) +EP_TEST_CONFIG = [ + (TOTAL_CARD_NUM, True, DEFAULT_HEIGHT, DEFAULT_WIDTH, True), +] + + +@pytest.mark.core_model +@pytest.mark.diffusion +@pytest.mark.parallel +@hardware_test(res={"cuda": "H100"}, num_cards={"cuda": TOTAL_CARD_NUM}) +@pytest.mark.parametrize("model_name", MODELS) +def test_ep(model_name): + baseline_cache: dict[str, InferenceResult] = {} + results: list[dict] = [] + # baseline_key = PROMPT + "_" + str(DEFAULT_HEIGHT) + "*" + str(DEFAULT_WIDTH) + "_" + str(DEFAULT_SEED) + "_" + str(DEFAULT_STEPS) + # baseline_cache[baseline_key] = _run_inference( + # model_name=model_name, + # tensor_parallel_size=TOTAL_CARD_NUM, + # enable_expert_parallel=False, + # warmup=True, + # ) + device_count = current_omni_platform.get_device_count() + + print("\n" + "=" * 90) + print(f"Sequence Parallel Test - Model: {model_name}") + print(f"Available GPUs: {device_count}") + print("=" * 90) + + for tensor_parallel_size, enable_ep, height, width, is_perf_test in EP_TEST_CONFIG: + cache_key = PROMPT + "_" + str(height) + "*" + str(width) + "_" + str(DEFAULT_SEED) + "_" + str(DEFAULT_STEPS) + if cache_key not in baseline_cache: + print( + f"\n--- Running baseline{{ prompt: {PROMPT}, resolution: {height}x{width}, seed: {DEFAULT_SEED}, steps: {DEFAULT_STEPS} }}---" + ) + baseline = _run_inference( + model_name=model_name, + tensor_parallel_size=TOTAL_CARD_NUM, + enable_expert_parallel=False, + warmup=True, + ) + assert len(baseline.images) == 1 + baseline_cache[cache_key] = baseline + print(f"[baseline] {height}x{width}: {baseline.elapsed_ms:.0f}ms") + else: + baseline = baseline_cache[cache_key] + + print(f"\n--- Running test: {{ enable_ep: {enable_ep} }} ---") + infer_result = _run_inference( + model_name=model_name, + tensor_parallel_size=tensor_parallel_size, + enable_expert_parallel=enable_ep, + warmup=is_perf_test, + ) + assert len(infer_result.images) == 1 + mean_diff, max_diff, cos_sim_mean, cos_sim_max, mse = _diff_metrics(baseline.images[0], infer_result.images[0]) + print(f"{cos_sim_mean}, {cos_sim_max}, {mse}") + + # Build result entry + result = { + "prompt": PROMPT, + "enable_ep": enable_ep, + "height": height, + "width": width, + "baseline_ms": baseline.elapsed_ms, + "ep_ms": infer_result.elapsed_ms, + "mean_diff": mean_diff, + "max_diff": max_diff, + "cos_sim_mean": cos_sim_mean, + "cos_sim_max": cos_sim_max, + "mse": mse, + "is_perf_test": is_perf_test, + } + results.append(result) + + # Output based on test type + if is_perf_test: + speedup = baseline.elapsed_ms / infer_result.elapsed_ms if infer_result.elapsed_ms > 0 else 0 + result["speedup"] = speedup + print( + f"[enable_ep: {enable_ep}] {tensor_parallel_size} GPUs | " + f"baseline: {baseline.elapsed_ms:.0f}ms, ep: {infer_result.elapsed_ms:.0f}ms, " + f"speedup: {speedup:.2f}x" + ) + else: + print( + f"[enable_ep: {enable_ep}] {tensor_parallel_size} GPUs | ep: {infer_result.elapsed_ms:.0f}ms (correctness only)" + ) + + print( + f"[enable_ep: {enable_ep}] diff: [mean={mean_diff:.6e}, max={max_diff:.6e}], cos_sim: [mean={cos_sim_mean:.6e}, max={cos_sim_max:.6e}], mse: {mse:.6e}" + ) + assert mean_diff <= DIFF_MEAN_THRESHOLD and max_diff <= DIFF_MAX_THRESHOLD, ( + f"[enable_ep: {enable_ep}] output differs from baseline: mean={mean_diff:.6e}, max={max_diff:.6e}" + ) + + print("\n" + "=" * 90) + print("SUMMARY") + print("=" * 90) + print(f"{'Mode':<15} {'GPUs':<6} {'Size':<10} {'Baseline':<12} {'EP':<12} {'Speedup':<10} {'Status'}") + print("-" * 90) + for r in results: + speedup_str = f"{r['speedup']:.2f}x" if r.get("speedup") else "N/A" + baseline_str = f"{r['baseline_ms']:.0f}ms" if r["is_perf_test"] else "N/A" + status = "PASS" if r["mean_diff"] <= DIFF_MEAN_THRESHOLD else "FAIL" + print( + f"{r['prompt']:<15.10} {r['enable_ep']:<6} {r['height']}x{r['width']:<5} " + f"{baseline_str:<12} {r['ep_ms']:.0f}ms{'':<7} {speedup_str:<10} {status}" + ) + print("=" * 90) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 6a1e44472..38366469e 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -43,6 +43,9 @@ class DiffusionParallelConfig: tensor_parallel_size: int = 1 """Number of tensor parallel groups.""" + enable_expert_parallel: bool = False + """Enable expert parallelism for MoE layers (TP is still used for non-MoE layers).""" + sequence_parallel_size: int | None = None """Number of sequence parallel groups. sequence_parallel_size = ring_degree * ulysses_degree""" @@ -457,6 +460,17 @@ class OmniDiffusionConfig: quantization: str | None = None quantization_config: "DiffusionQuantizationConfig | dict[str, Any] | None" = None + @property + def is_moe(self) -> bool: + num_experts = self.tf_model_config.get("num_experts", None) + if isinstance(num_experts, int): + return num_experts > 1 + + if isinstance(num_experts, (list, tuple)): + return any(isinstance(n, int) and n > 1 for n in num_experts) + + return False + def settle_port(self, port: int, port_inc: int = 42, max_attempts: int = 100) -> int: """ Find an available port with retry logic. diff --git a/vllm_omni/diffusion/distributed/parallel_state.py b/vllm_omni/diffusion/distributed/parallel_state.py index 4c09fd7a0..82c9c1a45 100644 --- a/vllm_omni/diffusion/distributed/parallel_state.py +++ b/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() diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 19ba5f31c..1b9dc5667 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -106,6 +106,7 @@ class DiffusionWorker: vllm_config = VllmConfig(compilation_config=CompilationConfig()) vllm_config.parallel_config.tensor_parallel_size = self.od_config.parallel_config.tensor_parallel_size vllm_config.parallel_config.data_parallel_size = self.od_config.parallel_config.data_parallel_size + vllm_config.parallel_config.enable_expert_parallel = self.od_config.parallel_config.enable_expert_parallel self.vllm_config = vllm_config # Initialize distributed environment @@ -127,6 +128,7 @@ class DiffusionWorker: pipeline_parallel_size=parallel_config.pipeline_parallel_size, fully_shard_degree=parallel_config.hsdp_shard_size if parallel_config.use_hsdp else 1, hsdp_replicate_size=parallel_config.hsdp_replicate_size if parallel_config.use_hsdp else 1, + enable_expert_parallel=parallel_config.enable_expert_parallel, ) init_workspace_manager(self.device) diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 6a4b88dbf..3c3cb4cd6 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -157,6 +157,7 @@ class AsyncOmni(OmniBase): ring_degree = kwargs.get("ring_degree") or 1 sequence_parallel_size = kwargs.get("sequence_parallel_size") tensor_parallel_size = kwargs.get("tensor_parallel_size") or 1 + enable_expert_parallel = kwargs.get("enable_expert_parallel") or False cfg_parallel_size = kwargs.get("cfg_parallel_size") or 1 vae_patch_parallel_size = kwargs.get("vae_patch_parallel_size") or 1 use_hsdp = kwargs.get("use_hsdp", False) @@ -179,6 +180,7 @@ class AsyncOmni(OmniBase): pipeline_parallel_size=1, data_parallel_size=1, tensor_parallel_size=tensor_parallel_size, + enable_expert_parallel=enable_expert_parallel, sequence_parallel_size=sequence_parallel_size, ulysses_degree=ulysses_degree, ring_degree=ring_degree, diff --git a/vllm_omni/entrypoints/async_omni_diffusion.py b/vllm_omni/entrypoints/async_omni_diffusion.py index 5da6e81ed..37d358d45 100644 --- a/vllm_omni/entrypoints/async_omni_diffusion.py +++ b/vllm_omni/entrypoints/async_omni_diffusion.py @@ -100,6 +100,7 @@ class AsyncOmniDiffusion: if cfg is None: raise ValueError(f"Could not find config.json or model_index.json for model {od_config.model}") + od_config.tf_model_config = TransformerConfig.from_dict(cfg) model_type = cfg.get("model_type") architectures = cfg.get("architectures") or [] # Bagel/NextStep models don't have a model_index.json, so we set the pipeline class name manually diff --git a/vllm_omni/entrypoints/omni_diffusion.py b/vllm_omni/entrypoints/omni_diffusion.py index 59ddf2eeb..fbfcec1d4 100644 --- a/vllm_omni/entrypoints/omni_diffusion.py +++ b/vllm_omni/entrypoints/omni_diffusion.py @@ -98,7 +98,7 @@ class OmniDiffusion: if od_config.model_class_name is None: od_config.model_class_name = pipeline_class - od_config.tf_model_config = TransformerConfig() + od_config.tf_model_config = TransformerConfig().from_dict(cfg) od_config.update_multimodal_support() if cfg_kv_collect_func is not None: