[Refactor][CI] Reorganize test tree, and move slow e2e to weekly (#4951)
Signed-off-by: wangyu <410167048@qq.com>
This commit is contained in:
wangyu authored and GitHub committed
1 parent
925347332f
commit
7c17111ff5
158 files changed
+3479
-4130
No files matched your search
@@ -1,151 +1,181 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""VAE patch-parallel decode parity tests for diffusion models.
|
||||
|
||||
Compares ``vae_patch_parallel_size=2`` against ``vae_patch_parallel_size=1`` on
|
||||
TP=2 with VAE tiling enabled. Covers one video model (Wan T2V) and one image
|
||||
model (Qwen-Image). CUDA/ROCm-only (>=2 devices).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from PIL import Image
|
||||
from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
|
||||
|
||||
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
|
||||
|
||||
# ruff: noqa: E402
|
||||
REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
if str(REPO_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
import time
|
||||
|
||||
from tests.helpers.runtime import OmniRunner
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
all_model_configs = [
|
||||
{
|
||||
"model_name": "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
|
||||
"out_width": 1280,
|
||||
"out_height": 720,
|
||||
"out_frames": 16,
|
||||
},
|
||||
{
|
||||
"model_name": "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
||||
"out_width": 1280,
|
||||
"out_height": 720,
|
||||
"out_frames": 16,
|
||||
},
|
||||
{
|
||||
"model_name": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
|
||||
"out_width": 1280,
|
||||
"out_height": 720,
|
||||
"out_frames": 16,
|
||||
},
|
||||
{
|
||||
"model_name": "Qwen/Qwen-Image",
|
||||
"out_width": 3072,
|
||||
"out_height": 3072,
|
||||
"out_frames": 1,
|
||||
},
|
||||
# {
|
||||
# "model_name": "stepfun-ai/NextStep-1.1",
|
||||
# "out_width": 3072,
|
||||
# "out_height": 3072,
|
||||
# "out_frames": 1,
|
||||
# },
|
||||
PROMPT = "A cat sitting on a table"
|
||||
|
||||
MODEL_CASES = [
|
||||
pytest.param(
|
||||
{
|
||||
"model_name": "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
||||
"height": 512,
|
||||
"width": 512,
|
||||
"num_frames": 5,
|
||||
"needs_image": False,
|
||||
"is_moe": False,
|
||||
"mean_threshold": 3e-2,
|
||||
"p99_threshold": 3e-2,
|
||||
},
|
||||
id="wan22_t2v_a14b",
|
||||
),
|
||||
pytest.param(
|
||||
{
|
||||
"model_name": "Qwen/Qwen-Image",
|
||||
"height": 1152,
|
||||
"width": 1152,
|
||||
"num_frames": 1,
|
||||
"needs_image": False,
|
||||
"is_moe": False,
|
||||
"mean_threshold": 5e-3,
|
||||
"p99_threshold": 1e-1,
|
||||
},
|
||||
id="qwen_image",
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def is_nextstep_model(model_name: str) -> bool:
|
||||
"""Check if the model is a NextStep model by reading its config."""
|
||||
from vllm.transformers_utils.config import get_hf_file_to_dict
|
||||
|
||||
try:
|
||||
cfg = get_hf_file_to_dict("config.json", model_name)
|
||||
if cfg and cfg.get("model_type") == "nextstep":
|
||||
return True
|
||||
except Exception:
|
||||
pass
|
||||
return False
|
||||
def _to_float_array(data: Any) -> np.ndarray:
|
||||
if isinstance(data, Image.Image):
|
||||
return np.asarray(data.convert("RGB"), dtype=np.float32) / 255.0
|
||||
if isinstance(data, torch.Tensor):
|
||||
arr = data.detach().cpu().numpy()
|
||||
else:
|
||||
arr = np.asarray(data)
|
||||
if arr.dtype == np.uint8:
|
||||
return arr.astype(np.float32) / 255.0
|
||||
return arr.astype(np.float32)
|
||||
|
||||
|
||||
def model_run(model_configs, tp, out_height, out_width, out_frames, using_tile, vae_patch_parallel_size=1):
|
||||
def _extract_output_array(output: Any) -> np.ndarray:
|
||||
req_out = output.request_output
|
||||
payload = req_out.images[0]
|
||||
return _to_float_array(payload)
|
||||
|
||||
|
||||
def _diff_metrics(a: np.ndarray, b: np.ndarray) -> tuple[float, float]:
|
||||
assert a.shape == b.shape, f"Output shapes differ: {a.shape} vs {b.shape}"
|
||||
abs_diff = np.abs(a - b)
|
||||
return float(abs_diff.mean()), float(np.quantile(abs_diff, 0.99))
|
||||
|
||||
|
||||
def _save_preview(arr: np.ndarray, path: Path) -> None:
|
||||
"""Save a single RGB frame preview; video tensors use the first frame."""
|
||||
preview = np.asarray(arr)
|
||||
# (B, T, H, W, C) video
|
||||
if preview.ndim == 5:
|
||||
preview = preview[0, 0]
|
||||
# (T, H, W, C) video or (B, H, W, C) image batch
|
||||
elif preview.ndim == 4:
|
||||
preview = preview[0]
|
||||
if preview.ndim != 3 or preview.shape[-1] not in (3, 4):
|
||||
raise ValueError(f"Expected HWC RGB(A) preview, got shape {arr.shape} -> {preview.shape}")
|
||||
preview = np.clip(preview[..., :3], 0.0, 1.0)
|
||||
Image.fromarray((preview * 255.0).astype(np.uint8)).save(path)
|
||||
|
||||
|
||||
def _build_sampling_params(model_case: dict[str, Any], seed: int) -> OmniDiffusionSamplingParams:
|
||||
kwargs: dict[str, Any] = {
|
||||
"height": model_case["height"],
|
||||
"width": model_case["width"],
|
||||
"num_frames": model_case["num_frames"],
|
||||
"num_inference_steps": 2,
|
||||
"seed": seed,
|
||||
}
|
||||
if model_case["model_name"].startswith("Wan-AI/"):
|
||||
kwargs["guidance_scale"] = 4.0
|
||||
if model_case["is_moe"]:
|
||||
kwargs["guidance_scale_2"] = 1.0
|
||||
kwargs["boundary_ratio"] = 0.5
|
||||
return OmniDiffusionSamplingParams(**kwargs)
|
||||
|
||||
|
||||
def _run_generate(
|
||||
*,
|
||||
model_case: dict[str, Any],
|
||||
vae_patch_parallel_size: int,
|
||||
seed: int,
|
||||
) -> np.ndarray:
|
||||
current_omni_platform.empty_cache()
|
||||
parallel_config = DiffusionParallelConfig(
|
||||
tensor_parallel_size=tp,
|
||||
tensor_parallel_size=2,
|
||||
vae_patch_parallel_size=vae_patch_parallel_size,
|
||||
)
|
||||
|
||||
omni_kwargs = {
|
||||
"vae_use_tiling": using_tile,
|
||||
"parallel_config": parallel_config,
|
||||
}
|
||||
use_nextstep = is_nextstep_model(model_configs["model_name"])
|
||||
if use_nextstep:
|
||||
# NextStep-1.1 requires explicit pipeline class
|
||||
omni_kwargs["model_class_name"] = "NextStep11Pipeline"
|
||||
|
||||
with OmniRunner(model_configs["model_name"], **omni_kwargs) as runner:
|
||||
m = runner.omni
|
||||
image = Image.new("RGB", (out_width, out_height), (0, 0, 0))
|
||||
start = time.perf_counter()
|
||||
outputs = m.generate(
|
||||
{
|
||||
"prompt": "A cat sitting on a table",
|
||||
"multi_modal_data": {"image": image},
|
||||
},
|
||||
sampling_params_list=OmniDiffusionSamplingParams(
|
||||
height=model_configs["out_height"],
|
||||
width=model_configs["out_width"],
|
||||
num_frames=model_configs["out_frames"],
|
||||
num_inference_steps=2,
|
||||
generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
|
||||
),
|
||||
with OmniRunner(
|
||||
model_case["model_name"],
|
||||
parallel_config=parallel_config,
|
||||
vae_use_tiling=True,
|
||||
enforce_eager=False,
|
||||
) as runner:
|
||||
request: dict[str, Any] = {"prompt": model_case.get("prompt", PROMPT)}
|
||||
if model_case["needs_image"]:
|
||||
request["multi_modal_data"] = {
|
||||
"image": Image.new("RGB", (model_case["width"], model_case["height"]), (0, 0, 0))
|
||||
}
|
||||
outputs = runner.omni.generate(
|
||||
request,
|
||||
sampling_params_list=_build_sampling_params(model_case, seed),
|
||||
)
|
||||
end = time.perf_counter()
|
||||
first_output = outputs[0]
|
||||
req_out = first_output.request_output
|
||||
frames = req_out.images[0]
|
||||
if isinstance(frames, torch.Tensor):
|
||||
frames = frames.detach().cpu().numpy()
|
||||
elif isinstance(frames, Image.Image):
|
||||
frames = np.array(frames)
|
||||
# frames shape: (batch, num_frames, height, width, channels)
|
||||
cost = (end - start) * 1000
|
||||
return frames, cost
|
||||
cleanup_dist_env_and_memory()
|
||||
return _extract_output_array(outputs[0])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_configs", all_model_configs)
|
||||
def test_vae_parallel_model(model_configs: str):
|
||||
parallel_result, parallel_time = model_run(
|
||||
model_configs=model_configs,
|
||||
tp=2,
|
||||
out_width=1280,
|
||||
out_height=704,
|
||||
out_frames=5,
|
||||
using_tile=True,
|
||||
vae_patch_parallel_size=2,
|
||||
)
|
||||
non_parallel_result, non_parallel_time = model_run(
|
||||
model_configs=model_configs,
|
||||
tp=2,
|
||||
out_width=1280,
|
||||
out_height=704,
|
||||
out_frames=5,
|
||||
using_tile=True,
|
||||
vae_patch_parallel_size=1,
|
||||
)
|
||||
result_diff = np.abs(non_parallel_result - parallel_result)
|
||||
@pytest.mark.full_model
|
||||
@pytest.mark.diffusion
|
||||
@pytest.mark.parallel
|
||||
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})
|
||||
@pytest.mark.parametrize("model_case", MODEL_CASES)
|
||||
def test_vae_patch_parallel_tp2(model_case: dict[str, Any], tmp_path: Path):
|
||||
if current_omni_platform.is_npu():
|
||||
pytest.skip("VAE patch parallel test is only supported on CUDA and ROCm for now.")
|
||||
if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2:
|
||||
pytest.skip("VAE patch parallel TP=2 requires >= 2 devices.")
|
||||
|
||||
mean_threshold = 3e-2
|
||||
max_threshold = 3e-2 # they should be totally same
|
||||
seed = 42
|
||||
model_id = model_case["model_name"].split("/")[-1]
|
||||
|
||||
baseline = _run_generate(model_case=model_case, vae_patch_parallel_size=1, seed=seed)
|
||||
parallel = _run_generate(model_case=model_case, vae_patch_parallel_size=2, seed=seed)
|
||||
|
||||
baseline_path = tmp_path / f"{model_id}_tp2_vae_pp1.png"
|
||||
parallel_path = tmp_path / f"{model_id}_tp2_vae_pp2.png"
|
||||
_save_preview(baseline, baseline_path)
|
||||
_save_preview(parallel, parallel_path)
|
||||
|
||||
mean_abs_diff, p99_abs_diff = _diff_metrics(baseline, parallel)
|
||||
mean_threshold = model_case["mean_threshold"]
|
||||
p99_threshold = model_case["p99_threshold"]
|
||||
print(
|
||||
f"{model_configs['model_name']} TP = 2 (tile + parallel vs tile): "
|
||||
f"mean_abs_diff={result_diff.mean():.6e}, max_abs_diff={result_diff.max():.6e}; "
|
||||
f"thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}; "
|
||||
f"parallel generate take time: {parallel_time:.2f} ms, non-parallel take time: {non_parallel_time:.2f} ms"
|
||||
f"{model_case['model_name']} VAE patch parallel diff (TP=2, pp=1 vs pp=2): "
|
||||
f"mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e}; "
|
||||
f"thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e}; "
|
||||
f"pp1_preview={baseline_path}, pp2_preview={parallel_path}"
|
||||
)
|
||||
assert mean_abs_diff <= mean_threshold and p99_abs_diff <= p99_threshold, (
|
||||
f"Output diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, "
|
||||
f"p99_abs_diff={p99_abs_diff:.6e} "
|
||||
f"(thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e})"
|
||||
)
|
||||
# assert non_parallel_time > parallel_time
|
||||
assert result_diff.mean() < mean_threshold
|
||||
assert result_diff.max() < max_threshold
|
||||
Reference in new issue
Block a user