wehub-resource-sync
498 lines
17 KiB
Python
498 lines
17 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""
|
|
System test for Sequence Parallel (SP) backends: Ulysses and Ring attention.
|
|
|
|
Tests verify that SP inference produces correct outputs compared to baseline.
|
|
"""
|
|
|
|
import gc
|
|
import os
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
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.diffusion.distributed.utils import build_local_sp_padding_mask
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
from vllm_omni.platforms import current_omni_platform
|
|
|
|
# ruff: noqa: E402
|
|
REPO_ROOT = Path(__file__).resolve().parents[3]
|
|
if str(REPO_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(REPO_ROOT))
|
|
|
|
|
|
# Test configuration
|
|
MODELS = ["riverclouds/qwen_image_random"]
|
|
PROMPT = "a photo of a cat sitting on a laptop keyboard"
|
|
DEFAULT_HEIGHT = 256
|
|
DEFAULT_WIDTH = 256
|
|
DEFAULT_SEED = 42
|
|
DEFAULT_STEPS = 4
|
|
DIFF_MEAN_THRESHOLD = 2e-2
|
|
DIFF_MAX_THRESHOLD = 2e-1
|
|
|
|
|
|
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]:
|
|
"""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)
|
|
return abs_diff.mean().item(), abs_diff.max().item()
|
|
|
|
|
|
def _run_inference(
|
|
model_name: str,
|
|
dtype: torch.dtype,
|
|
attn_backend: str,
|
|
ulysses_degree: int = 1,
|
|
ring_degree: int = 1,
|
|
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(ulysses_degree=ulysses_degree, ring_degree=ring_degree)
|
|
try:
|
|
with OmniRunner(
|
|
model_name,
|
|
parallel_config=parallel_config,
|
|
dtype=dtype,
|
|
attention_backend=attn_backend,
|
|
) 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=0.0,
|
|
generator=torch.Generator(current_omni_platform.device_type).manual_seed(seed + 1000),
|
|
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=0.0,
|
|
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.images,
|
|
elapsed_ms=elapsed_ms,
|
|
)
|
|
finally:
|
|
_cleanup_distributed()
|
|
|
|
|
|
# =============================================================================
|
|
# Correctness & Performance Tests
|
|
# =============================================================================
|
|
|
|
# SP configurations: (ulysses_degree, ring_degree, height, width, warmup, is_perf_test)
|
|
# - warmup: whether to run warmup for this SP config
|
|
# - is_perf_test: whether this is a performance test (show speedup metrics)
|
|
SP_CONFIGS_L2 = [
|
|
# Hybrid - correctness only
|
|
(2, 2, DEFAULT_HEIGHT, DEFAULT_WIDTH, False, False),
|
|
]
|
|
|
|
SP_CONFIGS_L3 = [
|
|
# Ulysses-2 - performance test
|
|
(2, 1, DEFAULT_HEIGHT, DEFAULT_WIDTH, True, True),
|
|
(1, 2, DEFAULT_HEIGHT, DEFAULT_WIDTH, True, True), # Ring-2 - performance test
|
|
# Hybrid - correctness only
|
|
(2, 2, DEFAULT_HEIGHT, DEFAULT_WIDTH, False, False),
|
|
(4, 1, 272, 272, False, False), # Ulysses-4 - shape and correctness
|
|
]
|
|
|
|
|
|
def _get_sp_mode(ulysses_degree: int, ring_degree: int) -> str:
|
|
"""Get SP mode name for logging."""
|
|
if ulysses_degree > 1 and ring_degree == 1:
|
|
return f"ulysses-{ulysses_degree}"
|
|
elif ring_degree > 1 and ulysses_degree == 1:
|
|
return f"ring-{ring_degree}"
|
|
else:
|
|
return f"hybrid-{ulysses_degree}x{ring_degree}"
|
|
|
|
|
|
@pytest.mark.core_model
|
|
@pytest.mark.diffusion
|
|
@pytest.mark.parallel
|
|
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 2, "rocm": 2})
|
|
@pytest.mark.parametrize("model_name", MODELS)
|
|
def test_sp_correctness(model_name: str):
|
|
"""Test that SP inference produces correct outputs and measure performance.
|
|
|
|
Runs baseline once per unique (height, width), then tests all SP configs.
|
|
|
|
Note: Run with `pytest -v -s` to see detailed output.
|
|
"""
|
|
device_count = current_omni_platform.get_device_count()
|
|
|
|
# Cache baseline results by (height, width)
|
|
# Key: (height, width), Value: (result, warmup_used)
|
|
baseline_cache: dict[tuple[int, int], InferenceResult] = {}
|
|
|
|
# Collect results for summary
|
|
results: list[dict] = []
|
|
|
|
print("\n" + "=" * 70)
|
|
print(f"Sequence Parallel Test - Model: {model_name}")
|
|
print(f"Available GPUs: {device_count}")
|
|
print("=" * 70)
|
|
|
|
for ulysses_degree, ring_degree, height, width, sp_warmup, is_perf_test in SP_CONFIGS_L2:
|
|
sp_size = ulysses_degree * ring_degree
|
|
sp_mode = _get_sp_mode(ulysses_degree, ring_degree)
|
|
|
|
if device_count < sp_size:
|
|
print(f"\n[{sp_mode}] SKIPPED (requires {sp_size} GPUs)")
|
|
continue
|
|
|
|
# Determine baseline warmup: only for default size (performance tests)
|
|
cache_key = (height, width)
|
|
baseline_warmup = height == DEFAULT_HEIGHT and width == DEFAULT_WIDTH
|
|
|
|
# Get or compute baseline for this (height, width)
|
|
if cache_key not in baseline_cache:
|
|
print(f"\n--- Running baseline {height}x{width} (warmup={baseline_warmup}) ---")
|
|
baseline = _run_inference(
|
|
model_name,
|
|
torch.bfloat16,
|
|
"sdpa",
|
|
height=height,
|
|
width=width,
|
|
warmup=baseline_warmup,
|
|
)
|
|
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]
|
|
|
|
# Run SP
|
|
print(f"\n--- Running {sp_mode} (warmup={sp_warmup}) ---")
|
|
sp_result = _run_inference(
|
|
model_name,
|
|
torch.bfloat16,
|
|
"sdpa",
|
|
ulysses_degree=ulysses_degree,
|
|
ring_degree=ring_degree,
|
|
height=height,
|
|
width=width,
|
|
warmup=sp_warmup,
|
|
)
|
|
assert len(sp_result.images) == 1
|
|
|
|
# Compare outputs (correctness)
|
|
mean_diff, max_diff = _diff_metrics(baseline.images[0], sp_result.images[0])
|
|
|
|
# Build result entry
|
|
result = {
|
|
"mode": sp_mode,
|
|
"sp_size": sp_size,
|
|
"height": height,
|
|
"width": width,
|
|
"baseline_ms": baseline.elapsed_ms,
|
|
"sp_ms": sp_result.elapsed_ms,
|
|
"mean_diff": mean_diff,
|
|
"max_diff": max_diff,
|
|
"is_perf_test": is_perf_test,
|
|
}
|
|
results.append(result)
|
|
|
|
# Output based on test type
|
|
if is_perf_test:
|
|
speedup = baseline.elapsed_ms / sp_result.elapsed_ms if sp_result.elapsed_ms > 0 else 0
|
|
result["speedup"] = speedup
|
|
print(
|
|
f"[{sp_mode}] {sp_size} GPUs | "
|
|
f"baseline: {baseline.elapsed_ms:.0f}ms, sp: {sp_result.elapsed_ms:.0f}ms, "
|
|
f"speedup: {speedup:.2f}x"
|
|
)
|
|
else:
|
|
print(f"[{sp_mode}] {sp_size} GPUs | sp: {sp_result.elapsed_ms:.0f}ms (correctness only)")
|
|
|
|
print(f"[{sp_mode}] diff: mean={mean_diff:.6e}, max={max_diff:.6e}")
|
|
|
|
# Assert correctness
|
|
assert mean_diff <= DIFF_MEAN_THRESHOLD and max_diff <= DIFF_MAX_THRESHOLD, (
|
|
f"[{sp_mode}] SP output differs from baseline: mean={mean_diff:.6e}, max={max_diff:.6e}"
|
|
)
|
|
|
|
# Summary
|
|
print("\n" + "=" * 70)
|
|
print("SUMMARY")
|
|
print("=" * 70)
|
|
print(f"{'Mode':<15} {'GPUs':<6} {'Size':<10} {'Baseline':<12} {'SP':<12} {'Speedup':<10} {'Status'}")
|
|
print("-" * 70)
|
|
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['mode']:<15} {r['sp_size']:<6} {r['height']}x{r['width']:<5} "
|
|
f"{baseline_str:<12} {r['sp_ms']:.0f}ms{'':<7} {speedup_str:<10} {status}"
|
|
)
|
|
print("=" * 70)
|
|
|
|
|
|
# TODO: After PR#1272 is merged, add markers
|
|
# @pytest.mark.advanced_model
|
|
@pytest.mark.diffusion
|
|
@pytest.mark.parallel
|
|
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})
|
|
@pytest.mark.parametrize("model_name", MODELS)
|
|
def test_sp_correctness_advanced(model_name: str):
|
|
"""Test that SP inference produces correct outputs and measure performance.
|
|
|
|
Runs baseline once per unique (height, width), then tests all SP configs.
|
|
|
|
Note: Run with `pytest -v -s` to see detailed output.
|
|
"""
|
|
device_count = current_omni_platform.get_device_count()
|
|
|
|
# Cache baseline results by (height, width)
|
|
# Key: (height, width), Value: (result, warmup_used)
|
|
baseline_cache: dict[tuple[int, int], InferenceResult] = {}
|
|
|
|
# Collect results for summary
|
|
results: list[dict] = []
|
|
|
|
print("\n" + "=" * 70)
|
|
print(f"Sequence Parallel Test - Model: {model_name}")
|
|
print(f"Available GPUs: {device_count}")
|
|
print("=" * 70)
|
|
|
|
for ulysses_degree, ring_degree, height, width, sp_warmup, is_perf_test in SP_CONFIGS_L3:
|
|
sp_size = ulysses_degree * ring_degree
|
|
sp_mode = _get_sp_mode(ulysses_degree, ring_degree)
|
|
|
|
if device_count < sp_size:
|
|
print(f"\n[{sp_mode}] SKIPPED (requires {sp_size} GPUs)")
|
|
continue
|
|
|
|
# Determine baseline warmup: only for default size (performance tests)
|
|
cache_key = (height, width)
|
|
baseline_warmup = height == DEFAULT_HEIGHT and width == DEFAULT_WIDTH
|
|
|
|
# Get or compute baseline for this (height, width)
|
|
if cache_key not in baseline_cache:
|
|
print(f"\n--- Running baseline {height}x{width} (warmup={baseline_warmup}) ---")
|
|
baseline = _run_inference(
|
|
model_name,
|
|
torch.bfloat16,
|
|
"sdpa",
|
|
height=height,
|
|
width=width,
|
|
warmup=baseline_warmup,
|
|
)
|
|
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]
|
|
|
|
# Run SP
|
|
print(f"\n--- Running {sp_mode} (warmup={sp_warmup}) ---")
|
|
sp_result = _run_inference(
|
|
model_name,
|
|
torch.bfloat16,
|
|
"sdpa",
|
|
ulysses_degree=ulysses_degree,
|
|
ring_degree=ring_degree,
|
|
height=height,
|
|
width=width,
|
|
warmup=sp_warmup,
|
|
)
|
|
assert len(sp_result.images) == 1
|
|
|
|
# Compare outputs (correctness)
|
|
mean_diff, max_diff = _diff_metrics(baseline.images[0], sp_result.images[0])
|
|
|
|
# Build result entry
|
|
result = {
|
|
"mode": sp_mode,
|
|
"sp_size": sp_size,
|
|
"height": height,
|
|
"width": width,
|
|
"baseline_ms": baseline.elapsed_ms,
|
|
"sp_ms": sp_result.elapsed_ms,
|
|
"mean_diff": mean_diff,
|
|
"max_diff": max_diff,
|
|
"is_perf_test": is_perf_test,
|
|
}
|
|
results.append(result)
|
|
|
|
# Output based on test type
|
|
if is_perf_test:
|
|
speedup = baseline.elapsed_ms / sp_result.elapsed_ms if sp_result.elapsed_ms > 0 else 0
|
|
result["speedup"] = speedup
|
|
print(
|
|
f"[{sp_mode}] {sp_size} GPUs | "
|
|
f"baseline: {baseline.elapsed_ms:.0f}ms, sp: {sp_result.elapsed_ms:.0f}ms, "
|
|
f"speedup: {speedup:.2f}x"
|
|
)
|
|
else:
|
|
print(f"[{sp_mode}] {sp_size} GPUs | sp: {sp_result.elapsed_ms:.0f}ms (correctness only)")
|
|
|
|
print(f"[{sp_mode}] diff: mean={mean_diff:.6e}, max={max_diff:.6e}")
|
|
|
|
# Assert correctness
|
|
assert mean_diff <= DIFF_MEAN_THRESHOLD and max_diff <= DIFF_MAX_THRESHOLD, (
|
|
f"[{sp_mode}] SP output differs from baseline: mean={mean_diff:.6e}, max={max_diff:.6e}"
|
|
)
|
|
|
|
# Summary
|
|
print("\n" + "=" * 70)
|
|
print("SUMMARY")
|
|
print("=" * 70)
|
|
print(f"{'Mode':<15} {'GPUs':<6} {'Size':<10} {'Baseline':<12} {'SP':<12} {'Speedup':<10} {'Status'}")
|
|
print("-" * 70)
|
|
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['mode']:<15} {r['sp_size']:<6} {r['height']}x{r['width']:<5} "
|
|
f"{baseline_str:<12} {r['sp_ms']:.0f}ms{'':<7} {speedup_str:<10} {status}"
|
|
)
|
|
print("=" * 70)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not (current_omni_platform.is_cuda() or current_omni_platform.is_xpu()),
|
|
reason="Only tested on CUDA and XPU",
|
|
)
|
|
@pytest.mark.diffusion
|
|
@pytest.mark.parallel
|
|
@pytest.mark.core_model
|
|
def test_local_sp_padding_mask(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""A partially padded SP shard must receive a local-length mask."""
|
|
mask = build_local_sp_padding_mask(
|
|
batch_size=2,
|
|
local_seq_len=4,
|
|
sp_original_seq_len=5,
|
|
sp_padding_size=3,
|
|
sequence_parallel_rank=1,
|
|
device=torch.device(current_omni_platform.device_type),
|
|
)
|
|
|
|
expected = torch.tensor(
|
|
[
|
|
[True, False, False, False],
|
|
[True, False, False, False],
|
|
],
|
|
dtype=torch.bool,
|
|
device=mask.device,
|
|
)
|
|
assert mask is not None
|
|
assert mask.shape == (2, 4)
|
|
assert torch.equal(mask, expected)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not (current_omni_platform.is_cuda() or current_omni_platform.is_xpu()),
|
|
reason="Only tested on CUDA and XPU",
|
|
)
|
|
@pytest.mark.diffusion
|
|
@pytest.mark.parallel
|
|
@pytest.mark.core_model
|
|
def test_local_sp_padding_mask_no_padding(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""A rank whose local shard contains no padding should not get a mask."""
|
|
mask = build_local_sp_padding_mask(
|
|
batch_size=2,
|
|
local_seq_len=4,
|
|
sp_original_seq_len=5,
|
|
sp_padding_size=3,
|
|
sequence_parallel_rank=0,
|
|
device=torch.device(current_omni_platform.device_type),
|
|
)
|
|
|
|
assert mask is None
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not (current_omni_platform.is_cuda() or current_omni_platform.is_xpu()),
|
|
reason="Only tested on CUDA and XPU",
|
|
)
|
|
@pytest.mark.diffusion
|
|
@pytest.mark.parallel
|
|
@pytest.mark.core_model
|
|
def test_wan_sp_plan() -> None:
|
|
"""Wan2.2 must shard hidden states before CacheDiT-wrapped transformer blocks."""
|
|
try:
|
|
from vllm_omni.diffusion.distributed.sp_plan import validate_sp_plan
|
|
from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3DModel
|
|
except ImportError as exc:
|
|
pytest.skip(f"WanTransformer3DModel not available: {exc}")
|
|
|
|
plan = getattr(WanTransformer3DModel, "_sp_plan", None)
|
|
|
|
assert plan is not None
|
|
assert "_sp_shard_point" in plan
|
|
assert "blocks.0" not in plan
|
|
|
|
shard_plan = plan["_sp_shard_point"]
|
|
assert 0 in shard_plan
|
|
assert shard_plan[0].split_dim == 1
|
|
assert shard_plan[0].expected_dims == 3
|
|
assert shard_plan[0].split_output is True
|
|
|
|
validate_sp_plan(plan)
|