diff --git a/tests/diffusion/distributed/test_wan_spatial_shard.py b/tests/diffusion/distributed/test_wan_spatial_shard.py index 72a2a45f5..beb5ac2db 100644 --- a/tests/diffusion/distributed/test_wan_spatial_shard.py +++ b/tests/diffusion/distributed/test_wan_spatial_shard.py @@ -231,6 +231,128 @@ def test_reshard_from_trimmed_width_pads_invalid_columns(monkeypatch: pytest.Mon assert torch.equal(out.flatten(), torch.tensor([4.0, 0.0])) +@pytest.mark.core_model +@pytest.mark.cpu +@pytest.mark.parametrize("split_dim", ["height", "width"]) +@pytest.mark.parametrize( + ("full_extent", "world_size"), + [ + pytest.param(5, 4, id="four-ranks-empty-tail"), + pytest.param(20, 8, id="production-eight-ranks-empty-tail"), + pytest.param(39, 2, id="two-ranks-partial-tail"), + ], +) +def test_reshard_from_trimmed_extent_clamps_and_reconstructs( + monkeypatch: pytest.MonkeyPatch, + split_dim: str, + full_extent: int, + world_size: int, +): + local_extent = (full_extent + world_size - 1) // world_size + shape = (1, 1, 1, full_extent, 1) if split_dim == "height" else (1, 1, 1, 1, full_extent) + x = torch.arange(full_extent, dtype=torch.float32).reshape(shape) + outputs = [] + + for rank in range(world_size): + monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group, rank=rank: (rank, world_size)) + token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set( + wan_spatial_shard.SpatialShardContext( + input_extent=full_extent, + local_input_extent=local_extent, + split_dim=split_dim, + rank=rank, + world_size=world_size, + ) + ) + try: + out = wan_spatial_shard.reshard_from_trimmed_extent( + x, + local_extent=local_extent, + split_dim=split_dim, + group=object(), + ) + finally: + wan_spatial_shard._SPATIAL_SHARD_CONTEXT.reset(token) + + start = rank * local_extent + valid_extent = min(local_extent, max(0, full_extent - start)) + expected = torch.zeros(local_extent, dtype=torch.float32) + if valid_extent: + expected[:valid_extent] = torch.arange(start, start + valid_extent, dtype=torch.float32) + assert out.shape[_spatial_dim_for_test(split_dim)] == local_extent + assert torch.equal(out.flatten(), expected) + outputs.append(out) + + reconstructed = torch.cat(outputs, dim=_spatial_dim_for_test(split_dim)) + reconstructed = reconstructed.narrow(_spatial_dim_for_test(split_dim), 0, full_extent) + assert torch.equal(reconstructed, x) + + +def _spatial_dim_for_test(split_dim: str) -> int: + return -2 if split_dim == "height" else -1 + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_reshard_from_trimmed_extent_rejects_rank_mismatch(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (1, 2)) + x = torch.zeros((1, 1, 1, 1, 4), dtype=torch.float32) + token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set( + wan_spatial_shard.SpatialShardContext( + input_extent=4, + local_input_extent=2, + split_dim="width", + rank=0, + world_size=2, + ) + ) + try: + with pytest.raises(RuntimeError, match="group_rank=1, context_rank=0"): + wan_spatial_shard.reshard_from_trimmed_extent( + x, + local_extent=2, + split_dim="width", + group=object(), + ) + finally: + wan_spatial_shard._SPATIAL_SHARD_CONTEXT.reset(token) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_attention_wrapper_rejects_spatial_extent_change(monkeypatch: pytest.MonkeyPatch): + class ShrinkingWanAttentionBlock(torch.nn.Module): + def forward(self, x): + return x[..., :-1] + + monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (0, 2)) + monkeypatch.setattr( + wan_spatial_shard, + "all_gather_along_dim", + 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") + x = torch.zeros((1, 1, 1, 1, 2), dtype=torch.float32) + token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set( + wan_spatial_shard.SpatialShardContext( + input_extent=4, + local_input_extent=2, + split_dim="width", + rank=0, + world_size=2, + ) + ) + try: + with pytest.raises( + RuntimeError, + match=r"local_extent=2, gathered_extent=4, trimmed_extent=4, output_extent=3", + ): + module(x) + finally: + wan_spatial_shard._SPATIAL_SHARD_CONTEXT.reset(token) + + @pytest.mark.core_model @pytest.mark.cpu def test_halo_exchange_single_rank_noop(monkeypatch: pytest.MonkeyPatch): diff --git a/vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py b/vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py index bda728540..07edad694 100644 --- a/vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py +++ b/vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py @@ -203,8 +203,17 @@ def reshard_from_trimmed_extent( return x dim = _spatial_dim(split_dim) + ctx = _SPATIAL_SHARD_CONTEXT.get() + if ctx is not None and ctx.rank != rank: + raise RuntimeError( + "Wan VAE spatial-shard rank mismatch while resharding: " + f"group_rank={rank}, context_rank={ctx.rank}, split_dim={split_dim!r}." + ) + valid_extent = _local_valid_extent(local_extent) - start = rank * local_extent + actual_extent = x.shape[dim] + start = min(rank * local_extent, actual_extent) + valid_extent = min(valid_extent, actual_extent - start) local = _narrow_along_dim(x, dim, start, valid_extent).contiguous() if valid_extent < local_extent: local = _pad_along_dim(local, local_extent - valid_extent, dim=dim) @@ -594,10 +603,22 @@ def _patch_attention_block(module: nn.Module, group: dist.ProcessGroup, split_di dim = _spatial_dim(split_dim) local_extent = x.shape[dim] gathered = all_gather_along_dim(x, group=group, dim=dim).contiguous() + gathered_extent = gathered.shape[dim] full_extent = _current_full_extent(local_extent) if full_extent is not None: - gathered = _narrow_along_dim(gathered, dim, 0, full_extent).contiguous() + trim_extent = min(full_extent, gathered_extent) + gathered = _narrow_along_dim(gathered, dim, 0, trim_extent).contiguous() + else: + trim_extent = gathered_extent out = orig_forward(gathered, *args, **kwargs) + if out.shape[dim] != trim_extent: + raise RuntimeError( + "Wan VAE attention changed the spatial extent during global gather: " + f"rank={_rank_world(group)[0]}, split_dim={split_dim!r}, " + f"local_extent={local_extent}, gathered_extent={gathered_extent}, " + f"trimmed_extent={trim_extent}, output_extent={out.shape[dim]}, " + f"context={_SPATIAL_SHARD_CONTEXT.get()!r}." + ) return reshard_from_trimmed_extent(out, local_extent=local_extent, split_dim=split_dim, group=group) module.forward = MethodType(_forward, module)