diff --git a/tests/diffusion/distributed/test_wan_spatial_shard.py b/tests/diffusion/distributed/test_wan_spatial_shard.py index 2990cba14..b89e32907 100644 --- a/tests/diffusion/distributed/test_wan_spatial_shard.py +++ b/tests/diffusion/distributed/test_wan_spatial_shard.py @@ -11,7 +11,6 @@ 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 @@ -352,7 +351,7 @@ def test_attention_wrapper_rejects_spatial_extent_change(monkeypatch: pytest.Mon lambda x, *, group, dim, dst=None: torch.cat([x, x], dim=dim), ) module = ShrinkingWanAttentionBlock() - wan_spatial_shard._patch_attention_block(module, group=object(), split_dim="width") + wan_spatial_shard._patch_attention_block(module, group=object()) x = torch.zeros((1, 1, 1, 1, 2), dtype=torch.float32) token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set( wan_spatial_shard.SpatialShardContext( @@ -490,23 +489,18 @@ def test_auto_spatial_shard_gate_requires_full_group(monkeypatch: pytest.MonkeyP @pytest.mark.core_model @pytest.mark.cpu -def test_spatial_wrappers_are_prepared_only_for_a_full_group(monkeypatch: pytest.MonkeyPatch): +def test_set_parallel_size_prepares_spatial_wrappers_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"), - ) + executor = SimpleNamespace(group=object(), parallel_size=1, parallel_mode="tile") + + def set_parallel_size(parallel_size: int, mode: str = "tile") -> None: + executor.parallel_size = parallel_size + executor.parallel_mode = mode + + executor.set_parallel_size = set_parallel_size + object.__setattr__(vae, "distributed_executor", executor) 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, @@ -517,11 +511,10 @@ def test_spatial_wrappers_are_prepared_only_for_a_full_group(monkeypatch: pytest lambda candidate, group, split_dim: installed.append((candidate, group, split_dim)), ) - prepare_pipeline_wan_spatial_shard_decode(torch.nn.Identity()) + vae.set_parallel_size(2, mode="auto") assert installed == [] - vae.distributed_executor.parallel_size = 4 - prepare_pipeline_wan_spatial_shard_decode(torch.nn.Identity()) + vae.set_parallel_size(4, mode="auto") assert installed == [(vae, vae.distributed_executor.group, "height")] diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py index 5447a8fc8..665196ad9 100644 --- a/tests/diffusion/test_diffusion_model_runner.py +++ b/tests/diffusion/test_diffusion_model_runner.py @@ -732,19 +732,6 @@ def test_execute_model_runs_forward_after_kv_receive(monkeypatch): @pytest.mark.core_model @pytest.mark.cpu def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch): - memory_pool_active = False - preparation_pool_states = [] - - @contextmanager - def _memory_pool_context(tag): - nonlocal memory_pool_active - assert tag == "weights" - memory_pool_active = True - try: - yield - finally: - memory_pool_active = False - class _DummyLoader: def __init__(self, load_config, od_config=None): del load_config, od_config @@ -792,22 +779,16 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch): monkeypatch.setattr(model_runner_module, "LoadConfig", lambda: object()) monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _DummyLoader) monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _DummyMemoryProfiler) - monkeypatch.setattr( - model_runner_module, - "prepare_pipeline_wan_spatial_shard_decode", - lambda pipeline: preparation_pool_states.append(memory_pool_active), - ) monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda od_config, device: None) monkeypatch.setattr( model_runner_module, "get_cache_backend", lambda cache_backend, cache_config: dummy_cache_backend ) - DiffusionModelRunner.load_model(runner, memory_pool_context_fn=_memory_pool_context) + DiffusionModelRunner.load_model(runner) assert runner.cache_backend is None assert runner.od_config.cache_backend is None assert dummy_cache_backend.enabled is False - assert preparation_pool_states == [True] @pytest.mark.core_model diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index c54c0e1e3..df10245c1 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -471,10 +471,10 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig): raise ValueError("allgather_degree > 1 is mutually exclusive with ulysses_degree/ring_degree > 1") if self.ulysses_mode not in {"strict", "advanced_uaa"}: raise ValueError("ulysses_mode must be 'strict' or 'advanced_uaa'") - if self.vae_parallel_mode not in {"auto", "tile", "spatial_shard_height", "spatial_shard_width"}: + if self.vae_parallel_mode not in {"tile", "auto", "spatial_shard_height", "spatial_shard_width"}: raise ValueError( "vae_parallel_mode must be one of " - "{'auto', 'tile', 'spatial_shard_height', 'spatial_shard_width'}, " + "{'tile', 'auto', 'spatial_shard_height', 'spatial_shard_width'}, " f"but got {self.vae_parallel_mode!r}." ) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 2f1a4ede4..bc430a18c 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -210,11 +210,11 @@ class DiffusionParallelConfig: vae_parallel_mode: str = "tile" """VAE parallel decode strategy. + - "tile": Patch/tile parallel decode (default). Each rank decodes a subset + of spatial tiles and the results are stitched on rank 0. - "auto": Select the longer spatial-shard axis for validated Wan tiled decodes on a full DiT-group topology; otherwise use patch/tile parallel decode. - - "tile": Patch/tile parallel decode (default). Each rank decodes a subset - of spatial tiles and the results are stitched on rank 0. - "spatial_shard_height": Spatially-sharded decode that splits decoder feature maps along height and exchanges halo rows around spatial convolutions. @@ -254,9 +254,9 @@ class DiffusionParallelConfig: assert self.allgather_degree > 0, "AllGather degree must be > 0" assert self.cfg_parallel_size > 0, "CFG parallel size must be > 0" assert self.vae_patch_parallel_size > 0, "VAE patch parallel size must be > 0" - assert self.vae_parallel_mode in {"auto", "tile", "spatial_shard_height", "spatial_shard_width"}, ( + assert self.vae_parallel_mode in {"tile", "auto", "spatial_shard_height", "spatial_shard_width"}, ( "vae_parallel_mode must be one of " - "{'auto', 'tile', 'spatial_shard_height', 'spatial_shard_width'}, " + "{'tile', 'auto', 'spatial_shard_height', 'spatial_shard_width'}, " f"but got {self.vae_parallel_mode!r}." ) if self.allgather_degree > 1: diff --git a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py index 4b8f0b965..5290145a1 100644 --- a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py +++ b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py @@ -47,27 +47,6 @@ def select_auto_spatial_shard_split_dim( return max(candidates, key=lambda candidate: (candidate[0], candidate[1] == "width"))[1] -def prepare_pipeline_wan_spatial_shard_decode(pipeline: torch.nn.Module) -> None: - """Install eligible Wan wrappers while weight allocations are still tagged.""" - from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery - - if not dist.is_initialized(): - return - for vae in ModuleDiscovery.discover(pipeline).vaes: - if not isinstance(vae, DistributedAutoencoderKLWan): - continue - executor = vae.distributed_executor - mode = executor.parallel_mode - if mode not in {"auto", "spatial_shard_height", "spatial_shard_width"}: - continue - requested_size = int(executor.parallel_size) - world_size = dist.get_world_size(group=executor.group) - if requested_size <= 1 or requested_size != world_size: - continue - split_dim = "width" if mode == "spatial_shard_width" else "height" - wan_spatial_shard.install_wan_spatial_shard_decode(vae, executor.group, split_dim=split_dim) - - class OmniAutoencoderKLWan(AutoencoderKLWan): def _execution_context(self): try: @@ -101,6 +80,19 @@ class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin): model.init_distributed() return model + def set_parallel_size(self, parallel_size: int, mode: str = "tile") -> None: + super().set_parallel_size(parallel_size, mode=mode) + if mode not in {"auto", "spatial_shard_height", "spatial_shard_width"}: + return + + executor = self.distributed_executor + world_size = dist.get_world_size(group=executor.group) + if parallel_size <= 1 or parallel_size != world_size: + return + + split_dim = "width" if mode == "spatial_shard_width" else "height" + wan_spatial_shard.install_wan_spatial_shard_decode(self, executor.group, split_dim=split_dim) + def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]: _, _, num_frames, height, width = z.shape sample_height = height * self.spatial_compression_ratio diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 727bdf705..d118a9833 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -32,9 +32,6 @@ from vllm_omni.diffusion.compile import regionally_compile from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata -from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import ( - prepare_pipeline_wan_spatial_shard_decode, -) from vllm_omni.diffusion.forward_context import set_forward_context from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.interface import ( @@ -269,7 +266,6 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): custom_pipeline_name=custom_pipeline_name, device=self.device, ) - prepare_pipeline_wan_spatial_shard_decode(self.pipeline) time_after_load = time.perf_counter() logger.info( diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 87fd0b4b1..0350d5450 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -784,7 +784,7 @@ class OmniServeCommand(CLISubcommand): "--vae-parallel-mode", type=str, default="tile", - choices=["auto", "tile", "spatial_shard_height", "spatial_shard_width"], + choices=["tile", "auto", "spatial_shard_height", "spatial_shard_width"], help="VAE parallel decode strategy for diffusion models. " "'tile' (default) uses patch/tile decode; 'auto' selects the longer spatial-shard axis for " "validated Wan tiled decodes on a full DiT group, then falls back to patch/tile decode; "