# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni 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.helpers.mark import hardware_test from tests.helpers.runtime import OmniRunner 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, ) try: with OmniRunner(model_name, parallel_config=parallel_config) as runner: omni = runner.omni # 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].images, elapsed_ms=elapsed_ms, ) finally: _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): device_count = current_omni_platform.get_device_count() if device_count < TOTAL_CARD_NUM: pytest.skip(f"Requires at least {TOTAL_CARD_NUM} GPUs, got {device_count}") 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, # ) 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)