119 lines
3.9 KiB
Python
119 lines
3.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""Kandinsky 6 VAE patch parallel keeps the single-GPU tile grid."""
|
|
|
|
import os
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
import torch.distributed as dist
|
|
import torch.multiprocessing as mp
|
|
|
|
from vllm_omni.diffusion.models.kandinsky6.modeling_kandinsky6_vae import AutoencoderKLHunyuanVideo
|
|
|
|
|
|
def _stub_vae() -> AutoencoderKLHunyuanVideo:
|
|
vae = AutoencoderKLHunyuanVideo.__new__(AutoencoderKLHunyuanVideo)
|
|
vae.spatial_compression_ratio = 8
|
|
vae.temporal_compression_ratio = 4
|
|
vae.use_tiling = True
|
|
vae.use_framewise_decoding = True
|
|
vae.tile_sample_min_height = 32
|
|
vae.tile_sample_min_width = 32
|
|
vae.tile_sample_min_num_frames = 16
|
|
vae.tile_sample_stride_height = 24
|
|
vae.tile_sample_stride_width = 24
|
|
vae.tile_sample_stride_num_frames = 12
|
|
vae.tile_size = None
|
|
vae.distributed_executor = None
|
|
vae.post_quant_conv = lambda tensor: tensor
|
|
vae._seen_widths = []
|
|
|
|
def decoder(tensor: torch.Tensor) -> torch.Tensor:
|
|
vae._seen_widths.append(int(tensor.shape[-1]))
|
|
return tensor + 1
|
|
|
|
vae.decoder = decoder
|
|
return vae
|
|
|
|
|
|
@pytest.mark.core_model
|
|
@pytest.mark.cpu
|
|
def test_distributed_decode_tiles_match_local_tiled_decode():
|
|
vae = _stub_vae()
|
|
latent = torch.arange(1 * 4 * 2 * 8 * 12, dtype=torch.float32).reshape(1, 4, 2, 8, 12)
|
|
|
|
local = vae.tiled_decode(latent).sample
|
|
tasks, spec = vae._decode_tile_split(latent)
|
|
assert len(tasks) > 1
|
|
assert max(int(task.tensor.shape[-1]) for task in tasks) < latent.shape[-1]
|
|
|
|
decoded = {task.grid_coord: vae._decode_tile_exec(task) for task in tasks}
|
|
merged = vae._decode_tile_merge(decoded, spec)
|
|
|
|
torch.testing.assert_close(merged, local)
|
|
|
|
|
|
@pytest.mark.core_model
|
|
@pytest.mark.cpu
|
|
def test_decode_still_applies_optimal_tiling_when_patch_parallel_is_on():
|
|
vae = _stub_vae()
|
|
seen = {}
|
|
|
|
def tiling(shape, device=None):
|
|
del device
|
|
seen["shape"] = tuple(shape)
|
|
return (1, 17, 32, 32), (8, 24, 24)
|
|
|
|
from diffusers.models.autoencoders.vae import DecoderOutput
|
|
|
|
vae.get_dec_optimal_tiling = tiling
|
|
vae._decode = lambda z, return_dict=True: DecoderOutput(sample=z)
|
|
vae.distributed_executor = SimpleNamespace(parallel_size=2)
|
|
vae.is_distributed_enabled = lambda: True
|
|
|
|
out = vae.decode(torch.zeros(1, 16, 8, 8, 12))
|
|
|
|
assert seen["shape"] == (1, 16, 8, 8, 12)
|
|
assert tuple(out.sample.shape) == (1, 16, 8, 8, 12)
|
|
|
|
|
|
def _patch_parallel_worker(rank: int, port: str, shapes: dict) -> None:
|
|
os.environ["MASTER_ADDR"] = "127.0.0.1"
|
|
os.environ["MASTER_PORT"] = port
|
|
os.environ["RANK"] = str(rank)
|
|
os.environ["WORLD_SIZE"] = "2"
|
|
dist.init_process_group("gloo", rank=rank, world_size=2)
|
|
try:
|
|
import vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor as executor_mod
|
|
|
|
world = SimpleNamespace(device_group=dist.group.WORLD, cpu_group=dist.group.WORLD)
|
|
executor_mod.get_world_group = lambda: world
|
|
|
|
vae = _stub_vae()
|
|
latent = torch.arange(1 * 4 * 2 * 8 * 12, dtype=torch.float32).reshape(1, 4, 2, 8, 12)
|
|
reference = vae.tiled_decode(latent).sample.clone()
|
|
vae._seen_widths.clear()
|
|
vae.set_parallel_size(2, mode="tile")
|
|
decoded = vae.tiled_decode(latent).sample
|
|
torch.testing.assert_close(decoded, reference)
|
|
if rank == 0:
|
|
shapes["widths"] = list(vae._seen_widths)
|
|
shapes["full_width"] = int(latent.shape[-1])
|
|
finally:
|
|
dist.destroy_process_group()
|
|
|
|
|
|
@pytest.mark.core_model
|
|
@pytest.mark.cpu
|
|
def test_patch_parallel_decode_splits_tiles_across_ranks():
|
|
manager = mp.Manager()
|
|
shapes = manager.dict()
|
|
port = str(29500 + (os.getpid() % 1000))
|
|
mp.spawn(_patch_parallel_worker, args=(port, shapes), nprocs=2, join=True)
|
|
|
|
assert shapes["widths"]
|
|
assert max(shapes["widths"]) < shapes["full_width"]
|