[FEATURE] Tile/Patch parallelism refactor for easily support other models (#1366)

Signed-off-by: dengyunyang <584797741@qq.com>

tests/diffusion/distributed/test_distributed_vae_executor.py

@@ -0,0 +1,165 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+import torch.distributed as dist
+
+from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
+    DistributedVaeExecutor,
+    DistributedVaeMixin,
+    GridSpec,
+    TileTask,
+)
+
+
+class E2EOperator:
+    """tiles with (2, 3) -- (H,W)"""
+
+    def split(self, z):
+        rows_num = 2
+        cols_num = 3
+
+        h_size = z.shape[0]
+        w_size = z.shape[1]
+
+        tasks = []
+        for i in range(rows_num):
+            for j in range(cols_num):
+                tasks.append(
+                    TileTask(
+                        tile_id=len(tasks),
+                        grid_coord=(i, j),
+                        tensor=z[
+                            ((i * h_size) // rows_num) : (((i + 1) * h_size) // rows_num),
+                            ((j * w_size) // cols_num) : (((j + 1) * w_size) // cols_num),
+                        ],
+                    )
+                )
+
+        grid_spec = GridSpec(
+            split_dims=(0, 1),
+            grid_shape=(2, 3),
+        )
+        return tasks, grid_spec
+
+    def exec(self, task: TileTask):
+        return torch.full_like(task.tensor, fill_value=task.tile_id)
+
+    def merge(self, coord_tensor_map, grid_spec):
+        tiles = []
+        for r in range(grid_spec.grid_shape[0]):
+            row_tiles = []
+            for c in range(grid_spec.grid_shape[1]):
+                coord = (r, c)
+                row_tiles.append(coord_tensor_map[coord])
+            tiles.append(torch.cat(row_tiles, dim=1))
+        return torch.cat(tiles, dim=0)
+
+
+class DummyMixin(DistributedVaeMixin):
+    def __init__(self):
+        self.use_tiling = True
+        self.distributed_decoder = MagicMock()
+        self.distributed_decoder.parallel_size = 2
+        self.distributed_decoder.group = None
+
+
+@pytest.fixture(autouse=True)
+def mock_dist():
+    with (
+        patch.object(dist, "get_world_size", return_value=2),
+        patch.object(dist, "get_rank", return_value=0),
+        patch.object(dist, "is_initialized", return_value=True),
+        patch.object(dist, "gather", return_value=None),
+        patch.object(dist, "broadcast", return_value=None),
+    ):
+        yield
+
+
+@pytest.fixture(autouse=True)
+def mock_dit_group():
+    with patch(
+        "vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor.get_dit_group",
+        new=MagicMock(return_value=None),
+    ):
+        yield
+
+
+@pytest.fixture(autouse=True)
+def mock_dist_vae_executor():
+    with (
+        patch.object(DistributedVaeExecutor, "gather_tensors", side_effect=lambda x: [x]),
+        patch.object(DistributedVaeExecutor, "broadcast_tensor", side_effect=lambda x: x),
+    ):
+        yield
+
+
+# ============================
+# Unitest
+# ============================
+
+
+def test_balance_tasks():
+    executor = DistributedVaeExecutor()
+    workloads = [2, 5, 13, 8, 2, 4]  # each is 17
+    tasks = [TileTask(0, (i,), torch.tensor([i]), workload=load) for i, load in enumerate(workloads)]
+    assigned = executor._balance_tasks(tasks, num_rank=2)
+    assert len(assigned) == 2
+    total_work = [sum(t.workload for t in group) for group in assigned]
+    assert total_work[0] == total_work[1]
+
+
+def test_compute_global_padding_shape():
+    executor = DistributedVaeExecutor()
+    executor.rank = 0
+
+    local_results = [(0, torch.zeros((2, 3))), (1, torch.zeros((4, 2)))]
+    shape = executor._compute_global_padding_shape(local_results, 2, "cpu")
+
+    assert shape == [2, 4, 3]
+
+
+def test_pack_and_unpack():
+    executor = DistributedVaeExecutor()
+    executor.world_size = 1
+
+    grid_spec = GridSpec(split_dims=(0, 1), grid_shape=(2, 2))
+
+    # ======================
+    # pack
+    # ======================
+    local_results = [(0, torch.tensor([[1, 2], [3, 4]]))]
+
+    global_shape = [3, 3, 3]  # (tiles, H, W)
+
+    tile_tensor, meta_tensor = executor._pack_local_tiles(
+        local_results, global_shape, grid_spec, device="cpu", dtype=torch.int64
+    )
+
+    # check pack
+    assert tile_tensor.shape == torch.Size(global_shape)
+    assert meta_tensor.shape == (global_shape[0], len(grid_spec.split_dims) + 1)
+    assert meta_tensor[0, 0] == 0
+    assert meta_tensor[0, 1] == 2
+    assert meta_tensor[0, 2] == 2
+
+    # ======================
+    # unpack
+    # ======================
+    meta_gather = [meta_tensor]
+    tile_gather = [tile_tensor]
+
+    tid_coord_map = {0: (0, 0)}
+
+    coord_tensor_map = executor._unpack_tiles(meta_gather, tile_gather, grid_spec, tid_coord_map)
+
+    # check unpack
+    assert torch.equal(coord_tensor_map[(0, 0)], torch.tensor([[1, 2], [3, 4]]))
+
+
+def test_is_distributed_enabled():
+    mixin = DummyMixin()
+    assert mixin.is_distributed_enabled() is True
+
+    mixin.use_tiling = False
+    assert mixin.is_distributed_enabled() is False

tests/e2e/offline_inference/test_vae_decode_parallelism.py

@@ -0,0 +1,154 @@
+import os
+import sys
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
+
+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 vllm_omni import Omni
+from vllm_omni.platforms import current_omni_platform
+
+# os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+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,
+    # },
+]
+
+
+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 model_run(model_configs, tp, out_height, out_width, out_frames, using_tile, vae_patch_parallel_size=1):
+    try:
+        parallel_config = DiffusionParallelConfig(
+            tensor_parallel_size=tp,
+            vae_patch_parallel_size=vae_patch_parallel_size,
+        )
+
+        omni_kwargs = {
+            "model": model_configs["model_name"],
+            "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"
+        m = Omni(**omni_kwargs)
+        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),
+            ),
+        )
+        end = time.perf_counter()
+        first_output = outputs[0]
+        req_out = first_output.request_output[0]
+        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
+    finally:
+        m.close()
+        cleanup_dist_env_and_memory()
+
+
+@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)
+
+    mean_threshold = 3e-2
+    max_threshold = 3e-2  # they should be totally same
+    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"
+    )
+    # assert non_parallel_time > parallel_time
+    assert result_diff.mean() < mean_threshold
+    assert result_diff.max() < max_threshold