diff --git a/tests/e2e/offline_inference/test_zimage_parallelism.py b/tests/e2e/offline_inference/test_zimage_parallelism.py index b11176070..279710195 100644 --- a/tests/e2e/offline_inference/test_zimage_parallelism.py +++ b/tests/e2e/offline_inference/test_zimage_parallelism.py @@ -17,14 +17,12 @@ import time from pathlib import Path import numpy as np - -# import pytest +import pytest import torch from PIL import Image from vllm.distributed.parallel_state import cleanup_dist_env_and_memory -# from tests.utils import DeviceMemoryMonitor, hardware_test -from tests.utils import DeviceMemoryMonitor +from tests.utils import DeviceMemoryMonitor, hardware_test from vllm_omni import Omni from vllm_omni.diffusion.data import DiffusionParallelConfig from vllm_omni.inputs.data import OmniDiffusionSamplingParams @@ -37,7 +35,6 @@ if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) -# os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1" os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" PROMPT = "a photo of a cat sitting on a laptop keyboard" @@ -157,124 +154,124 @@ def _run_zimage_generate( cleanup_dist_env_and_memory() -# @pytest.mark.advanced_model -# @pytest.mark.diffusion -# @pytest.mark.parallel -# @hardware_test(res={"cuda": "L4", "rocm": "MI325", "xpu": "B60"}, num_cards={"cuda": 4, "rocm": 2, "xpu": 2}) -# def test_zimage_tensor_parallel_tp2(tmp_path: Path): -# if current_omni_platform.is_npu() or current_omni_platform.is_rocm(): -# pytest.skip("Z-Image TP e2e test is only supported on CUDA for now.") -# if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2: -# pytest.skip("Z-Image TP=2 requires >= 2 devices.") +@pytest.mark.advanced_model +@pytest.mark.diffusion +@pytest.mark.parallel +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) +def test_zimage_tensor_parallel_tp2(tmp_path: Path): + if current_omni_platform.is_npu() or current_omni_platform.is_rocm(): + pytest.skip("Z-Image TP e2e test is only supported on CUDA for now.") + if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2: + pytest.skip("Z-Image TP=2 requires >= 2 devices.") -# enforce_eager = False + enforce_eager = False -# height = 512 -# width = 512 -# num_inference_steps = 2 -# seed = 42 + height = 512 + width = 512 + num_inference_steps = 2 + seed = 42 -# tp1_img, tp1_time_s, tp1_peak_mem = _run_zimage_generate( -# tp_size=1, -# height=height, -# width=width, -# num_inference_steps=num_inference_steps, -# seed=seed, -# enforce_eager=enforce_eager, -# ) -# tp2_img, tp2_time_s, tp2_peak_mem = _run_zimage_generate( -# tp_size=2, -# height=height, -# width=width, -# num_inference_steps=num_inference_steps, -# seed=seed, -# enforce_eager=enforce_eager, -# ) + tp1_img, tp1_time_s, tp1_peak_mem = _run_zimage_generate( + tp_size=1, + height=height, + width=width, + num_inference_steps=num_inference_steps, + seed=seed, + enforce_eager=enforce_eager, + ) + tp2_img, tp2_time_s, tp2_peak_mem = _run_zimage_generate( + tp_size=2, + height=height, + width=width, + num_inference_steps=num_inference_steps, + seed=seed, + enforce_eager=enforce_eager, + ) -# tp1_path = tmp_path / "zimage_tp1.png" -# tp2_path = tmp_path / "zimage_tp2.png" -# tp1_img.save(tp1_path) -# tp2_img.save(tp2_path) + tp1_path = tmp_path / "zimage_tp1.png" + tp2_path = tmp_path / "zimage_tp2.png" + tp1_img.save(tp1_path) + tp2_img.save(tp2_path) -# assert tp1_img.width == width and tp1_img.height == height -# assert tp2_img.width == width and tp2_img.height == height + assert tp1_img.width == width and tp1_img.height == height + assert tp2_img.width == width and tp2_img.height == height -# mean_abs_diff, p99_abs_diff = _diff_metrics(tp1_img, tp2_img) -# mean_threshold = 3e-2 -# p99_threshold = 2.5e-1 -# print( -# "Z-Image TP image diff stats (TP=1 vs TP=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"tp1_img={tp1_path}, tp2_img={tp2_path}" -# ) -# assert mean_abs_diff <= mean_threshold and p99_abs_diff <= p99_threshold, ( -# f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e} " -# f"(thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e})" -# ) + mean_abs_diff, p99_abs_diff = _diff_metrics(tp1_img, tp2_img) + mean_threshold = 3e-2 + p99_threshold = 2.5e-1 + print( + "Z-Image TP image diff stats (TP=1 vs TP=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"tp1_img={tp1_path}, tp2_img={tp2_path}" + ) + assert mean_abs_diff <= mean_threshold and p99_abs_diff <= p99_threshold, ( + f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e} " + f"(thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e})" + ) -# print(f"Z-Image TP perf (lower is better): tp1_time_s={tp1_time_s:.6f}, tp2_time_s={tp2_time_s:.6f}") -# assert tp2_time_s < tp1_time_s, f"Expected TP=2 to be faster than TP=1 (tp1={tp1_time_s}, tp2={tp2_time_s})" + print(f"Z-Image TP perf (lower is better): tp1_time_s={tp1_time_s:.6f}, tp2_time_s={tp2_time_s:.6f}") + assert tp2_time_s < tp1_time_s, f"Expected TP=2 to be faster than TP=1 (tp1={tp1_time_s}, tp2={tp2_time_s})" -# print(f"Z-Image TP peak memory (MB): tp1_peak_mem={tp1_peak_mem:.2f}, tp2_peak_mem={tp2_peak_mem:.2f}") -# assert tp2_peak_mem < tp1_peak_mem, ( -# f"Expected TP=2 to use less peak memory than TP=1 (tp1={tp1_peak_mem}, tp2={tp2_peak_mem})" -# ) + print(f"Z-Image TP peak memory (MB): tp1_peak_mem={tp1_peak_mem:.2f}, tp2_peak_mem={tp2_peak_mem:.2f}") + assert tp2_peak_mem < tp1_peak_mem, ( + f"Expected TP=2 to use less peak memory than TP=1 (tp1={tp1_peak_mem}, tp2={tp2_peak_mem})" + ) -# @pytest.mark.integration -# def test_zimage_vae_patch_parallel_tp2(tmp_path: Path): -# if current_omni_platform.is_npu() or current_omni_platform.is_rocm(): -# pytest.skip("Z-Image VAE patch parallel e2e test is only supported on CUDA for now.") -# if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2: -# pytest.skip("Z-Image VAE patch parallel TP=2 requires >= 2 devices.") +@pytest.mark.integration +def test_zimage_vae_patch_parallel_tp2(tmp_path: Path): + if current_omni_platform.is_npu() or current_omni_platform.is_rocm(): + pytest.skip("Z-Image VAE patch parallel e2e test is only supported on CUDA for now.") + if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2: + pytest.skip("Z-Image VAE patch parallel TP=2 requires >= 2 devices.") -# enforce_eager = False + enforce_eager = False -# # Use a larger image to ensure there are multiple VAE tiles. -# height = 1152 -# width = 1152 -# num_inference_steps = 2 -# seed = 42 + # Use a larger image to ensure there are multiple VAE tiles. + height = 1152 + width = 1152 + num_inference_steps = 2 + seed = 42 -# baseline_img, _baseline_time_s, _baseline_peak_mem = _run_zimage_generate( -# tp_size=2, -# height=height, -# width=width, -# num_inference_steps=num_inference_steps, -# seed=seed, -# enforce_eager=enforce_eager, -# vae_use_tiling=True, -# vae_patch_parallel_size=1, -# num_requests=2, -# ) -# pp2_img, _pp2_time_s, _pp2_peak_mem = _run_zimage_generate( -# tp_size=2, -# height=height, -# width=width, -# num_inference_steps=num_inference_steps, -# seed=seed, -# enforce_eager=enforce_eager, -# vae_use_tiling=True, -# vae_patch_parallel_size=2, -# num_requests=2, -# ) + baseline_img, _baseline_time_s, _baseline_peak_mem = _run_zimage_generate( + tp_size=2, + height=height, + width=width, + num_inference_steps=num_inference_steps, + seed=seed, + enforce_eager=enforce_eager, + vae_use_tiling=True, + vae_patch_parallel_size=1, + num_requests=2, + ) + pp2_img, _pp2_time_s, _pp2_peak_mem = _run_zimage_generate( + tp_size=2, + height=height, + width=width, + num_inference_steps=num_inference_steps, + seed=seed, + enforce_eager=enforce_eager, + vae_use_tiling=True, + vae_patch_parallel_size=2, + num_requests=2, + ) -# baseline_path = tmp_path / "zimage_tp2_vae_pp1.png" -# pp2_path = tmp_path / "zimage_tp2_vae_pp2.png" -# baseline_img.save(baseline_path) -# pp2_img.save(pp2_path) + baseline_path = tmp_path / "zimage_tp2_vae_pp1.png" + pp2_path = tmp_path / "zimage_tp2_vae_pp2.png" + baseline_img.save(baseline_path) + pp2_img.save(pp2_path) -# mean_abs_diff, p99_abs_diff = _diff_metrics(baseline_img, pp2_img) -# mean_threshold = 5e-3 -# p99_threshold = 1e-1 -# print( -# "Z-Image VAE patch parallel image diff stats (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_img={baseline_path}, pp2_img={pp2_path}" -# ) -# assert mean_abs_diff <= mean_threshold and p99_abs_diff <= p99_threshold, ( -# f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e} " -# f"(thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e})" -# ) + mean_abs_diff, p99_abs_diff = _diff_metrics(baseline_img, pp2_img) + mean_threshold = 5e-3 + p99_threshold = 1e-1 + print( + "Z-Image VAE patch parallel image diff stats (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_img={baseline_path}, pp2_img={pp2_path}" + ) + assert mean_abs_diff <= mean_threshold and p99_abs_diff <= p99_threshold, ( + f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e} " + f"(thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e})" + ) diff --git a/tests/test_diffusion_config_propagation.py b/tests/test_diffusion_config_propagation.py new file mode 100644 index 000000000..58eb6097c --- /dev/null +++ b/tests/test_diffusion_config_propagation.py @@ -0,0 +1,108 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests that parallel_config survives the create_default_diffusion roundtrip. + +Regression tests for https://github.com/vllm-project/vllm-omni/issues/1862 +""" + +from collections.abc import Mapping + +import torch + +from vllm_omni.config.stage_config import StageConfigFactory +from vllm_omni.diffusion.data import ( + DiffusionParallelConfig, + OmniDiffusionConfig, +) + + +def _roundtrip_diffusion_config(**kwargs) -> OmniDiffusionConfig: + """Simulate the real path: create_default_diffusion → OmniDiffusionConfig. + + Does NOT manually reconstruct parallel_config — relies on + OmniDiffusionConfig.__post_init__ to handle the dict, just like + the production code path does. + """ + stages = StageConfigFactory.create_default_diffusion(kwargs) + engine_args = dict(stages[0]["engine_args"]) + return OmniDiffusionConfig.from_kwargs(**engine_args) + + +class TestParallelConfigPropagation: + """Core regression tests: parallel_config must survive serialization.""" + + def test_tp2_roundtrip(self): + pc = DiffusionParallelConfig(tensor_parallel_size=2) + od = _roundtrip_diffusion_config(model="test-model", parallel_config=pc) + assert od.parallel_config.tensor_parallel_size == 2 + assert od.parallel_config.world_size == 2 + + def test_tp4_devices_and_config(self): + pc = DiffusionParallelConfig(tensor_parallel_size=4) + stages = StageConfigFactory.create_default_diffusion({"parallel_config": pc, "model": "x"}) + assert stages[0]["runtime"]["devices"] == "0,1,2,3" + + # Let __post_init__ reconstruct from dict (real code path) + ea = dict(stages[0]["engine_args"]) + od = OmniDiffusionConfig.from_kwargs(**ea) + assert od.parallel_config.tensor_parallel_size == 4 + assert od.parallel_config.world_size == 4 + + def test_sp_config_roundtrip(self): + pc = DiffusionParallelConfig( + tensor_parallel_size=2, + ulysses_degree=2, + ring_degree=1, + ) + od = _roundtrip_diffusion_config(model="x", parallel_config=pc) + assert od.parallel_config.ulysses_degree == 2 + assert od.parallel_config.ring_degree == 1 + + def test_cfg_parallel_roundtrip(self): + pc = DiffusionParallelConfig(cfg_parallel_size=2) + od = _roundtrip_diffusion_config(model="x", parallel_config=pc) + assert od.parallel_config.cfg_parallel_size == 2 + assert od.parallel_config.world_size == 2 + + def test_no_parallel_config_defaults_to_tp1(self): + od = _roundtrip_diffusion_config(model="x") + assert od.parallel_config.tensor_parallel_size == 1 + assert od.parallel_config.world_size == 1 + + def test_num_gpus_derived_from_world_size(self): + pc = DiffusionParallelConfig(tensor_parallel_size=2) + od = _roundtrip_diffusion_config(model="x", parallel_config=pc) + assert od.num_gpus == 2 + + +class TestCreateDefaultDiffusion: + """Verify engine_args structure from create_default_diffusion.""" + + def test_parallel_config_serialized_as_dict(self): + """The key fix: parallel_config must appear in engine_args as a dict.""" + pc = DiffusionParallelConfig(tensor_parallel_size=2) + stages = StageConfigFactory.create_default_diffusion({"model": "x", "parallel_config": pc}) + ea = stages[0]["engine_args"] + assert "parallel_config" in ea + assert isinstance(ea["parallel_config"], Mapping) + assert ea["parallel_config"]["tensor_parallel_size"] == 2 + + def test_dtype_serialized_as_string(self): + stages = StageConfigFactory.create_default_diffusion({"dtype": torch.float16, "model": "x"}) + assert stages[0]["engine_args"]["dtype"] == "torch.float16" + + def test_cache_backend_defaults_to_none(self): + stages = StageConfigFactory.create_default_diffusion({"model": "x"}) + assert stages[0]["engine_args"]["cache_backend"] == "none" + + def test_single_gpu_default_devices(self): + stages = StageConfigFactory.create_default_diffusion({"model": "x"}) + assert stages[0]["runtime"]["devices"] == "0" + + def test_extra_kwargs_forwarded(self): + stages = StageConfigFactory.create_default_diffusion( + {"model": "x", "enforce_eager": True, "lora_path": "/tmp/lora"} + ) + ea = stages[0]["engine_args"] + assert ea["enforce_eager"] is True + assert ea["lora_path"] == "/tmp/lora" diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index ad298aa9d..e457c3d99 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -11,7 +11,7 @@ Runtime parameters (gpu_memory_utilization, tp_size, etc.) come from CLI. from __future__ import annotations import re -from dataclasses import dataclass, field +from dataclasses import asdict, dataclass, field from enum import Enum from pathlib import Path from typing import Any @@ -279,13 +279,16 @@ class StageConfigFactory: for i in range(1, num_devices): devices += f",{i}" - # Collect engine args – skip non-serializable objects engine_args: dict[str, Any] = {} for key, value in kwargs.items(): if key in ("parallel_config",): continue engine_args[key] = value + # Serialize parallel_config as dict for OmegaConf compatibility + if "parallel_config" in kwargs: + engine_args["parallel_config"] = asdict(kwargs["parallel_config"]) + engine_args.setdefault("cache_backend", "none") engine_args["model_stage"] = "diffusion" diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 3e44f861b..b67736c81 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -514,12 +514,11 @@ class OmniDiffusionConfig: initial_master_port = (self.master_port or 30005) + random.randint(0, 100) self.master_port = self.settle_port(initial_master_port, 37) - # Convert parallel_config dict to DiffusionParallelConfig if needed - # This must be done before accessing parallel_config.world_size - if isinstance(self.parallel_config, dict): - self.parallel_config = DiffusionParallelConfig.from_dict(self.parallel_config) + # Convert parallel_config dict/DictConfig to DiffusionParallelConfig + # Use Mapping to handle both plain dicts and OmegaConf DictConfig + if isinstance(self.parallel_config, Mapping): + self.parallel_config = DiffusionParallelConfig.from_dict(dict(self.parallel_config)) elif not isinstance(self.parallel_config, DiffusionParallelConfig): - # If it's neither dict nor DiffusionParallelConfig, use default config self.parallel_config = DiffusionParallelConfig() if self.num_gpus is None: