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/diffusion/attention/test_attention_sp.py b/tests/diffusion/attention/test_attention_sp.py index 234ac1563..c99ba9772 100644 --- a/tests/diffusion/attention/test_attention_sp.py +++ b/tests/diffusion/attention/test_attention_sp.py @@ -1,5 +1,29 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Tests for Ulysses + Ring sequence-parallel attention correctness. + +What is tested +-------------- +* ``test_sequence_parallel`` verifies that the ``Attention`` layer produces + numerically equivalent results whether the sequence is processed on a single + rank (baseline, SP=1) or sharded across multiple ranks via Ulysses/Ring SP. + The test spawns two separate multi-process runs with ``torch.multiprocessing.spawn``: + 1. Baseline – world_size=1, ulysses_degree=1, ring_degree=1. + 2. SP run – world_size=ulysses_degree*ring_degree, each rank holds a + contiguous slice of the full sequence. + After both runs, rank-0 output tensors are compared element-wise with a + tolerance appropriate for the dtype (bfloat16). + +SP-plan hooks are NOT applied in this test +------------------------------------------ +``ForwardContext.sp_plan_hooks_applied`` remains ``False``. As a result, +``ForwardContext.sp_active`` falls back to the "no hooks" branch: SP is +considered active whenever ``parallel_config.sequence_parallel_size > 1``. +This makes the test self-contained and suitable for standalone CI runs that +do not exercise the full model-registry pipeline. +""" + import os import pickle import tempfile @@ -27,6 +51,11 @@ def update_environment_variables(envs_dict: dict[str, str]): os.environ[k] = v +def seed_everything(seed: int): + torch.manual_seed(seed) + current_omni_platform.manual_seed(seed) + + class TestAttentionModel(torch.nn.Module): """Test model using Attention layer.""" @@ -138,11 +167,10 @@ class TestMultiLayerAttentionModel(torch.nn.Module): @pytest.mark.parametrize("num_heads", [8]) @pytest.mark.parametrize("head_size", [8]) @pytest.mark.parametrize("causal", [False]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) # [torch.float16, torch.bfloat16] +@pytest.mark.parametrize("dtype", [torch.bfloat16]) @pytest.mark.parametrize("use_sync", [False]) @pytest.mark.parametrize("dynamic", [False]) @pytest.mark.parametrize("use_compile", [False]) -@pytest.mark.parametrize("attn_backend", ["sdpa", "flash_attn"]) def test_sequence_parallel( ulysses_degree: int, ring_degree: int, @@ -156,7 +184,6 @@ def test_sequence_parallel( seq_len: int, num_heads: int, head_size: int, - attn_backend: str, ): """Test Ulysses attention by comparing with and without SP enabled.""" sequence_parallel_size = ulysses_degree * ring_degree @@ -200,7 +227,6 @@ def test_sequence_parallel( model_state_file, input_data_file, True, # is_baseline - attn_backend, ), nprocs=1, ) @@ -228,7 +254,6 @@ def test_sequence_parallel( model_state_file, input_data_file, False, # is_baseline - attn_backend, ), nprocs=sequence_parallel_size, ) @@ -324,12 +349,11 @@ def ulysses_attention_on_test_model( model_state_file: str, input_data_file: str, is_baseline: bool, - attn_backend: str, ): """Run Ulysses attention test on a test model and save results for comparison.""" # Use fixed seed for reproducibility across baseline and SP runs RANDOM_SEED = 42 - current_omni_platform.seed_everything(RANDOM_SEED) + seed_everything(RANDOM_SEED) mode_str = "Baseline (no SP)" if is_baseline else f"SP (ulysses={ulysses_degree}, ring={ring_degree})" print(f"\n[{mode_str}] Rank {local_rank}/{world_size} - Random seed set to {RANDOM_SEED}") @@ -366,7 +390,6 @@ def ulysses_attention_on_test_model( model="test_model", dtype=dtype, parallel_config=parallel_config, - attention_backend=attn_backend, # Set the attention backend here ) # Initialize model parallel @@ -411,11 +434,6 @@ def ulysses_attention_on_test_model( with open(model_state_file, "wb") as f: pickle.dump(model_state, f) - # Generate and save full input data with fixed seed - # Reinitialize RNG to ensure reproducibility - torch.manual_seed(42) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(42) full_hidden_states = torch.randn( (batch_size, seq_len, hidden_size), dtype=dtype, 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/forward_context.py b/vllm_omni/diffusion/forward_context.py index d2e68c8be..db574db1d 100644 --- a/vllm_omni/diffusion/forward_context.py +++ b/vllm_omni/diffusion/forward_context.py @@ -31,15 +31,35 @@ class ForwardContext: # Original sequence length before padding (for removing padding in gather) sp_original_seq_len: int | None = None - # SP active scope tracking + # Set by registry when _sp_plan hooks are applied. + # When True, sp_active is determined by _sp_shard_depth (for _sp_plan hooks) + # When False, sp_active defaults to True when sequence_parallel_size > 1 (for manual SP, standalone tests, etc.) + sp_plan_hooks_applied: bool = False + # SP active scope tracking within the _sp_plan hook mechanism. # Tracks the depth of SP sharding - incremented on shard, decremented on gather # Used by attention layers to determine if SP communication should be enabled _sp_shard_depth: int = 0 @property def sp_active(self) -> bool: - """Returns True when inside an SP sharded region (between shard and gather).""" - return self._sp_shard_depth > 0 + """Returns True when SP attention parallelism should be enabled. + + - If _sp_plan hooks are applied: use _sp_shard_depth (0 = outside sharded region). + - If _sp_plan hooks are NOT applied: default to True when sequence_parallel_size > 1, + since _sp_shard_depth is only meaningful within the _sp_plan hook mechanism. + """ + if self.sp_plan_hooks_applied: + return self._sp_shard_depth > 0 + # No _sp_plan: assume SP active when configured (manual SP, standalone tests) + if self.omni_diffusion_config is None: + raise ValueError( + "omni_diffusion_config is not set when checking sp_active! " + "This usually means set_forward_context() was not called. " + "Please call with set_forward_context(omni_diffusion_config=...)." + ) + + sp_size = self.omni_diffusion_config.parallel_config.sequence_parallel_size + return sp_size is not None and sp_size > 1 def __post_init__(self): pass diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index fe874c9e2..1b48c9c7d 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -10,6 +10,7 @@ from vllm.model_executor.models.registry import _LazyRegisteredModel, _ModelRegi from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import DistributedVaeMixin from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig, get_sp_plan_from_model +from vllm_omni.diffusion.forward_context import get_forward_context from vllm_omni.diffusion.hooks.sequence_parallel import apply_sequence_parallel logger = init_logger(__name__) @@ -264,6 +265,11 @@ def _apply_sequence_parallel_if_enabled(model, od_config: OmniDiffusionConfig) - apply_sequence_parallel(transformer, sp_config, plan) applied_count += 1 + # update forward context sp_plan_hooks_applied + ctx = get_forward_context() + ctx.sp_plan_hooks_applied = applied_count > 0 + logger.debug(f"Setting sp_plan_hooks_applied={ctx.sp_plan_hooks_applied} in ``ForwardContext``!") + if applied_count == 0: logger.warning( f"Sequence parallelism is enabled (sp_size={sp_size}) but no transformer with _sp_plan found. " 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: