anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

feat(diffusion): select Wan VAE decode mode automatically

Browse Source 
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
Hongsheng Liu committed 2026-08-11 09:41:05 +08:00
1 parent c27623c29c
commit 810112189b
16 files changed
+613 -186
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/diffusion/distributed/test_wan_spatial_shard.py
+233 -6
View File
@@ -1,4 +1,5 @@
import os
from contextlib import contextmanager
from types import SimpleNamespace
import pytest
@@ -10,6 +11,8 @@ from tests.helpers.mark import hardware_test
from vllm_omni.diffusion.distributed.autoencoders import wan_spatial_shard
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
DistributedAutoencoderKLWan,
prepare_pipeline_wan_spatial_shard_decode,
select_auto_spatial_shard_split_dim,
)
from vllm_omni.platforms import current_omni_platform
@@ -17,6 +20,23 @@ from vllm_omni.platforms import current_omni_platform
# the end uses hardware_test (H100 x2) plus full_model / diffusion / parallel.
@contextmanager
def _active_spatial_shard(split_dim: str):
token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set(
wan_spatial_shard.SpatialShardContext(
input_extent=2,
local_input_extent=2,
split_dim=split_dim,
rank=0,
world_size=3,
)
)
try:
yield
finally:
wan_spatial_shard._SPATIAL_SHARD_CONTEXT.reset(token)
@pytest.mark.core_model
@pytest.mark.cpu
def test_split_for_parallel_decode_pads_uneven_height():
@@ -255,12 +275,14 @@ def test_dist_zero_pad_only_applies_global_height_edges(monkeypatch: pytest.Monk
monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (1, 3))
mid_rank_pad = wan_spatial_shard.WanDistZeroPad2d((0, 1, 1, 1), group=object())
mid = mid_rank_pad(x)
with _active_spatial_shard("height"):
mid = mid_rank_pad(x)
assert mid.shape == (1, 1, 2, 3)
monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (2, 3))
last_rank_pad = wan_spatial_shard.WanDistZeroPad2d((0, 1, 1, 1), group=object())
last = last_rank_pad(x)
with _active_spatial_shard("height"):
last = last_rank_pad(x)
assert last.shape == (1, 1, 3, 3)
@@ -270,13 +292,15 @@ def test_dist_zero_pad_only_applies_global_width_edges(monkeypatch: pytest.Monke
x = torch.ones((1, 1, 2, 2))
monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (1, 3))
mid_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object(), split_dim="width")
mid = mid_rank_pad(x)
mid_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object())
with _active_spatial_shard("width"):
mid = mid_rank_pad(x)
assert mid.shape == (1, 1, 2, 2)
monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (2, 3))
last_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object(), split_dim="width")
last = last_rank_pad(x)
last_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object())
with _active_spatial_shard("width"):
last = last_rank_pad(x)
assert last.shape == (1, 1, 2, 3)
@@ -299,6 +323,113 @@ def test_spatial_shard_height_gate_falls_back_for_partial_group(monkeypatch: pyt
assert vae._spatial_shard_decode_enabled(z) is False
@pytest.mark.core_model
@pytest.mark.cpu
@pytest.mark.parametrize(
"z_shape,world_size,expected",
[
((1, 16, 1, 60, 104), 4, "width"),
((1, 16, 5, 60, 104), 4, "width"),
((1, 16, 5, 104, 60), 4, "height"),
((1, 16, 5, 2, 2), 4, None),
((1, 16, 5, 64, 64), 4, "width"),
],
)
def test_auto_spatial_shard_selector_is_shape_aware(z_shape, world_size, expected):
assert select_auto_spatial_shard_split_dim(z_shape, world_size) == expected
@pytest.mark.core_model
@pytest.mark.cpu
def test_auto_spatial_shard_gate_requires_full_group(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(
"vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan.dist.get_world_size",
lambda group=None: 4,
)
vae = DistributedAutoencoderKLWan.__new__(DistributedAutoencoderKLWan)
vae.use_tiling = True
vae.distributed_executor = SimpleNamespace(group=object(), parallel_size=4, parallel_mode="auto")
vae.is_distributed_enabled = lambda: True
large = torch.zeros((1, 16, 5, 60, 104))
unshardable = torch.zeros((1, 16, 1, 2, 2))
assert vae._spatial_shard_decode_split_dim(large, 4) == "width"
assert vae._spatial_shard_decode_enabled(large) is True
assert vae._spatial_shard_decode_enabled(unshardable) is False
vae.distributed_executor.parallel_size = 2
assert vae._spatial_shard_decode_enabled(large) is False
@pytest.mark.core_model
@pytest.mark.cpu
def test_spatial_wrappers_are_prepared_only_for_a_full_group(monkeypatch: pytest.MonkeyPatch):
vae = DistributedAutoencoderKLWan.__new__(DistributedAutoencoderKLWan)
object.__setattr__(
vae,
"distributed_executor",
SimpleNamespace(group=object(), parallel_size=2, parallel_mode="auto"),
)
installed = []
monkeypatch.setattr(
"vllm_omni.diffusion.offloader.module_collector.ModuleDiscovery.discover",
lambda pipeline: SimpleNamespace(vaes=[vae]),
)
monkeypatch.setattr(
"vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan.dist.is_initialized",
lambda: True,
)
monkeypatch.setattr(
"vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan.dist.get_world_size",
lambda group=None: 4,
)
monkeypatch.setattr(
wan_spatial_shard,
"install_wan_spatial_shard_decode",
lambda candidate, group, split_dim: installed.append((candidate, group, split_dim)),
)
prepare_pipeline_wan_spatial_shard_decode(torch.nn.Identity())
assert installed == []
vae.distributed_executor.parallel_size = 4
prepare_pipeline_wan_spatial_shard_decode(torch.nn.Identity())
assert installed == [(vae, vae.distributed_executor.group, "height")]
@pytest.mark.core_model
@pytest.mark.cpu
@torch.no_grad()
def test_installed_spatial_decoder_keeps_direct_fallback_and_can_change_axis():
vae = DistributedAutoencoderKLWan(
base_dim=8,
decoder_base_dim=8,
z_dim=4,
dim_mult=[1, 2],
num_res_blocks=1,
attn_scales=[],
temperal_downsample=[False],
latents_mean=[0.0] * 4,
latents_std=[1.0] * 4,
scale_factor_temporal=1,
scale_factor_spatial=2,
).eval()
latents = torch.randn(1, 4, 1, 4, 5)
parameter_names = {name for name, _ in vae.named_parameters()}
state_dict = {name: value.clone() for name, value in vae.state_dict().items()}
reference = vae.decode(latents, return_dict=False)[0]
wan_spatial_shard.install_wan_spatial_shard_decode(vae, group=object(), split_dim="height")
assert any(getattr(module, "_vllm_omni_dynamic_spatial_shard_conv", False) for module in vae.decoder.modules())
wan_spatial_shard.install_wan_spatial_shard_decode(vae, group=object(), split_dim="width")
assert {name for name, _ in vae.named_parameters()} == parameter_names
vae.load_state_dict(state_dict, strict=True)
assert torch.equal(vae.decode(latents, return_dict=False)[0], reference)
@pytest.mark.core_model
@pytest.mark.cpu
def test_spatial_shard_width_gate_selects_width(monkeypatch: pytest.MonkeyPatch):
@@ -444,3 +575,99 @@ def test_spatial_shard_decode_matches_reference(split_dim: str):
assert mean_abs_diff <= _SPATIAL_SHARD_TOLERANCE, (
f"spatial_shard_{split_dim} mean_abs_diff {mean_abs_diff} exceeds {_SPATIAL_SHARD_TOLERANCE}"
)
def _auto_spatial_shard_decode_worker(rank: int, return_dict, master_port: str) -> None:
from vllm_omni.diffusion.distributed.parallel_state import (
destroy_model_parallel,
init_distributed_environment,
initialize_model_parallel,
)
os.environ["MASTER_ADDR"] = "127.0.0.1"
os.environ["MASTER_PORT"] = master_port
device = current_omni_platform.get_torch_device(rank)
current_omni_platform.set_device(device)
backend = current_omni_platform.dist_backend
init_distributed_environment(world_size=2, rank=rank, local_rank=rank, backend=backend)
initialize_model_parallel(sequence_parallel_size=2, ulysses_degree=2, backend=backend)
try:
vae = DistributedAutoencoderKLWan.from_pretrained(
_SPATIAL_SHARD_MODEL,
subfolder=_SPATIAL_SHARD_SUBFOLDER,
torch_dtype=torch.float32,
).to(device=device, dtype=torch.float32)
vae.eval()
generator = torch.Generator(device=device).manual_seed(1)
small = torch.randn((1, vae.config.z_dim, 1, 60, 104), generator=generator, device=device)
landscape = torch.randn((1, vae.config.z_dim, 2, 60, 104), generator=generator, device=device)
portrait = torch.randn((1, vae.config.z_dim, 2, 104, 60), generator=generator, device=device)
with torch.inference_mode():
vae.use_tiling = False
vae.set_parallel_size(1, mode="tile")
small_reference = vae.decode(small, return_dict=False)[0]
vae.use_tiling = True
vae.set_parallel_size(2, mode="tile")
tile_before_install = vae.decode(small, return_dict=False)[0]
vae.set_parallel_size(2, mode="auto")
assert vae._spatial_shard_decode_split_dim(small, 2) == "width"
small_sharded = vae.decode(small, return_dict=False)[0]
assert vae._spatial_shard_decode_split_dim(landscape, 2) == "width"
vae.decode(landscape, return_dict=False)
vae.set_parallel_size(2, mode="tile")
tile_after_install = vae.decode(small, return_dict=False)[0]
vae.use_tiling = False
vae.set_parallel_size(1, mode="tile")
portrait_reference = vae.decode(portrait, return_dict=False)[0]
vae.use_tiling = True
vae.set_parallel_size(2, mode="auto")
assert vae._spatial_shard_decode_split_dim(portrait, 2) == "height"
portrait_sharded = vae.decode(portrait, return_dict=False)[0]
assert all(
buffer is None
for module in vae.decoder.modules()
for name, buffer in module._buffers.items()
if name in {"_halo_recv_top_buf", "_halo_recv_bottom_buf"}
)
if rank == 0:
small_diff = (small_sharded.float() - small_reference.float()).abs()
portrait_diff = (portrait_sharded.float() - portrait_reference.float()).abs()
return_dict["tile_fallback_exact"] = torch.equal(tile_after_install, tile_before_install)
return_dict["small_max_abs_diff"] = small_diff.max().item()
return_dict["small_mean_abs_diff"] = small_diff.mean().item()
return_dict["portrait_max_abs_diff"] = portrait_diff.max().item()
return_dict["portrait_mean_abs_diff"] = portrait_diff.mean().item()
finally:
destroy_model_parallel()
if dist.is_initialized():
dist.destroy_process_group()
@pytest.mark.full_model
@pytest.mark.diffusion
@pytest.mark.parallel
@hardware_test(res={"cuda": "H100"}, num_cards=2)
def test_auto_spatial_shard_switches_shape_paths_safely():
manager = mp.get_context("spawn").Manager()
return_dict = manager.dict()
mp.spawn(
_auto_spatial_shard_decode_worker,
args=(return_dict, "29502"),
nprocs=2,
join=True,
)
assert return_dict.get("tile_fallback_exact") is True
assert return_dict["small_max_abs_diff"] <= _SPATIAL_SHARD_TOLERANCE
assert return_dict["small_mean_abs_diff"] <= _SPATIAL_SHARD_TOLERANCE
assert return_dict["portrait_max_abs_diff"] <= _SPATIAL_SHARD_TOLERANCE
assert return_dict["portrait_mean_abs_diff"] <= _SPATIAL_SHARD_TOLERANCE