fix: propagate parallel_config through create_default_diffusion (#1878)

Signed-off-by: lishunyang <lishunyang12@163.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

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})"
+    )

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"

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"
 

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: