[LingBot World] Shard the streaming VAE decode across the Ulysses ranks (#7651)

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

tests/diffusion/distributed/test_wan_spatial_shard.py

@@ -493,6 +493,84 @@ def test_dist_zero_pad_only_applies_global_width_edges(monkeypatch: pytest.Monke
     assert last.shape == (1, 1, 2, 3)
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+class _PlainDecoder(torch.nn.Module):
+    """A decoder with nothing to patch, so the install only wraps ``forward``."""
+
+    def __init__(self):
+        super().__init__()
+        self.calls: list[dict] = []
+
+    def forward(self, x, feat_cache=None, feat_idx=None, first_chunk=False):
+        self.calls.append({"feat_cache": feat_cache, "feat_idx": feat_idx, "first_chunk": first_chunk})
+        return x
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize("world_size, dst", [(4, -1), (4, 4), (4, 5), (1, 1), (4, 1.5)])
+def test_install_rejects_invalid_destination_before_mutating_decoder(monkeypatch, world_size, dst):
+    monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (0, world_size))
+    vae = SimpleNamespace(decoder=torch.nn.Identity())
+    original_forward = vae.decoder.forward
+
+    def unexpected_patch(*args, **kwargs):
+        pytest.fail("invalid destination must fail before patching modules")
+
+    monkeypatch.setattr(wan_spatial_shard, "_patch_decoder_modules", unexpected_patch)
+    with pytest.raises(ValueError, match="dst must be None or an integer"):
+        wan_spatial_shard.install_wan_spatial_shard_decode(vae, object(), dst=dst)
+    assert vae.decoder.forward == original_forward
+    assert not getattr(vae, "_vllm_omni_wan_spatial_shard_installed", False)
+
+
+def _install_on_plain_decoder(monkeypatch: pytest.MonkeyPatch, *, dst):
+    monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (1, 2))
+    monkeypatch.setattr(
+        wan_spatial_shard,
+        "split_for_parallel_decode",
+        lambda x, *, upsample_count, split_dim, group: (x, x.shape[wan_spatial_shard._spatial_dim(split_dim)]),
+    )
+    gathers: list[dict] = []
+
+    def gather(x, *, expected_extent, split_dim, group, dst=None):
+        gathers.append({"expected_extent": expected_extent, "split_dim": split_dim, "dst": dst})
+        return x
+
+    monkeypatch.setattr(wan_spatial_shard, "gather_and_trim_extent", gather)
+    vae = SimpleNamespace(decoder=_PlainDecoder())
+    wan_spatial_shard.install_wan_spatial_shard_decode(vae, object(), split_dim="width", dst=dst)
+    return vae, gathers
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize("dst", [0, 1, None])
+def test_install_passes_the_assembling_rank_through_to_the_gather(monkeypatch: pytest.MonkeyPatch, dst):
+    vae, gathers = _install_on_plain_decoder(monkeypatch, dst=dst)
+    cache = [None]
+
+    out = vae.decoder(torch.zeros(1, 3, 1, 4, 6), feat_cache=cache, feat_idx=[0], first_chunk=True)
+
+    assert out.shape == (1, 3, 1, 4, 6)
+    assert gathers == [{"expected_extent": 6, "split_dim": "width", "dst": dst}]
+    # The streaming arguments reach the wrapped decoder untouched.
+    assert vae.decoder.calls == [{"feat_cache": cache, "feat_idx": [0], "first_chunk": True}]
+    assert vae._vllm_omni_wan_spatial_shard_dst == dst
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_install_refuses_to_change_the_assembling_rank(monkeypatch: pytest.MonkeyPatch):
+    vae, _ = _install_on_plain_decoder(monkeypatch, dst=None)
+
+    # Same settings again: a no-op, as before.
+    wan_spatial_shard.install_wan_spatial_shard_decode(vae, object(), split_dim="width", dst=None)
+    with pytest.raises(ValueError, match="already patched to assemble on None"):
+        wan_spatial_shard.install_wan_spatial_shard_decode(vae, object(), split_dim="width", dst=0)
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_spatial_shard_height_gate_falls_back_for_partial_group(monkeypatch: pytest.MonkeyPatch):
@@ -626,6 +704,105 @@ def _spatial_shard_decode_worker(rank: int, split_dim: str, return_dict, master_
             dist.destroy_process_group()
 
 
+_STREAMING_SHARD_CHUNKS = 3
+
+
+def _streaming_shard_worker(rank: int, return_dict, master_port: str) -> None:
+    """Sharded streaming decode with every rank keeping the frame, against an unsharded streaming reference.
+
+    Two sessions are interleaved chunk by chunk, so the per-session temporal cache
+    and the halo exchange are exercised across successive calls, then one session
+    is reset and must reproduce its opening chunk. Every rank checks every chunk.
+    """
+    from vllm_omni.diffusion.distributed.parallel_state import (
+        destroy_model_parallel,
+        get_sp_group,
+        init_distributed_environment,
+        initialize_model_parallel,
+    )
+    from vllm_omni.experimental.ar_diffusion.streaming_decode import WanStreamingDecoder
+
+    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)
+    dtype = torch.float32
+    backend = current_omni_platform.dist_backend
+    init_distributed_environment(world_size=_SPATIAL_SHARD_WORLD_SIZE, rank=rank, local_rank=rank, backend=backend)
+    initialize_model_parallel(
+        sequence_parallel_size=_SPATIAL_SHARD_WORLD_SIZE, ulysses_degree=_SPATIAL_SHARD_WORLD_SIZE, backend=backend
+    )
+    # VLLM_OMNI_TEST_WAN_VAE may point at a local diffusers checkpoint (with or
+    # without a ``vae`` subfolder) for machines without hub access.
+    model = os.environ.get("VLLM_OMNI_TEST_WAN_VAE", _SPATIAL_SHARD_MODEL)
+    local_vae_only = os.path.isdir(model) and not os.path.isdir(os.path.join(model, "vae"))
+    load = {"torch_dtype": dtype} if local_vae_only else {"subfolder": _SPATIAL_SHARD_SUBFOLDER, "torch_dtype": dtype}
+    try:
+        reference_vae = DistributedAutoencoderKLWan.from_pretrained(model, **load).to(device=device, dtype=dtype).eval()
+        sharded_vae = DistributedAutoencoderKLWan.from_pretrained(model, **load).to(device=device, dtype=dtype).eval()
+        wan_spatial_shard.install_wan_spatial_shard_decode(
+            sharded_vae, get_sp_group().device_group, split_dim="width", dst=None
+        )
+        reference = WanStreamingDecoder(reference_vae, bytes_per_pixel_fp32=1.0)
+        sharded = WanStreamingDecoder(sharded_vae, bytes_per_pixel_fp32=1.0)
+
+        generator = torch.Generator(device=device).manual_seed(0)
+        # Width 104 latents = 832 px: not divisible by two ranks' halo-padded shards without padding.
+        shape = (1, reference_vae.config.z_dim, 1, _SPATIAL_SHARD_LATENT_HEIGHT, _SPATIAL_SHARD_LATENT_WIDTH)
+        chunks = {
+            name: [
+                torch.randn(shape, generator=generator, device=device, dtype=dtype)
+                for _ in range(_STREAMING_SHARD_CHUNKS)
+            ]
+            for name in ("a", "b")
+        }
+        max_diff = 0.0
+        with torch.inference_mode():
+            ref_states = {name: reference.new_decode_state(name) for name in chunks}
+            shard_states = {name: sharded.new_decode_state(name) for name in chunks}
+            for index in range(_STREAMING_SHARD_CHUNKS):
+                for name in ("a", "b"):
+                    expected = reference.decode_chunk(chunks[name][index], ref_states[name]).float()
+                    got = sharded.decode_chunk(chunks[name][index], shard_states[name]).float()
+                    if got.shape != expected.shape:
+                        raise AssertionError(
+                            f"rank {rank} {name}[{index}]: {tuple(got.shape)} != {tuple(expected.shape)}"
+                        )
+                    max_diff = max(max_diff, (got - expected).abs().max().item())
+            # Reset session a: a fresh state must reproduce its opening chunk, on this rank too.
+            shard_states["a"].release()
+            ref_states["a"].release()
+            shard_states["a"] = sharded.new_decode_state("a")
+            ref_states["a"] = reference.new_decode_state("a")
+            expected = reference.decode_chunk(chunks["a"][0], ref_states["a"]).float()
+            got = sharded.decode_chunk(chunks["a"][0], shard_states["a"]).float()
+            max_diff = max(max_diff, (got - expected).abs().max().item())
+        return_dict[f"max_abs_diff_rank{rank}"] = max_diff
+        return_dict[f"shape_rank{rank}"] = tuple(got.shape)
+    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", "B200"]}, num_cards=_SPATIAL_SHARD_WORLD_SIZE)
+def test_streaming_shard_decode_matches_unsharded_streaming_on_every_rank():
+    manager = mp.get_context("spawn").Manager()
+    return_dict = manager.dict()
+    mp.spawn(_streaming_shard_worker, args=(return_dict, "29510"), nprocs=_SPATIAL_SHARD_WORLD_SIZE, join=True)
+
+    for rank in range(_SPATIAL_SHARD_WORLD_SIZE):
+        assert f"max_abs_diff_rank{rank}" in return_dict, f"rank {rank} did not report"
+        diff = return_dict[f"max_abs_diff_rank{rank}"]
+        shape = return_dict[f"shape_rank{rank}"]
+        print(f"streaming shard rank {rank}: max_abs_diff={diff:.6e} shape={shape}")
+        assert shape[-1] == _SPATIAL_SHARD_LATENT_WIDTH * 8 and shape[-2] == _SPATIAL_SHARD_LATENT_HEIGHT * 8
+        assert diff <= _SPATIAL_SHARD_TOLERANCE, f"rank {rank}: {diff} exceeds {_SPATIAL_SHARD_TOLERANCE}"
+
+
 @pytest.mark.full_model
 @pytest.mark.diffusion
 @pytest.mark.parallel

tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py

@@ -737,6 +737,99 @@ def test_component_discovery_uses_official_checkpoint_contract() -> None:
     assert module._loader_state.prefetch_calls == [("checkpoint", ("tokenizer", "text_encoder", "vae"), False)]
 
 
+def _ulysses_config(size: int):
+    parallel_config = _od_config().parallel_config
+    parallel_config.sequence_parallel_size = size
+    parallel_config.ulysses_degree = size
+    return _od_config(parallel_config=parallel_config)
+
+
+def _record_shard_install(monkeypatch, module, *, world_size: int):
+    """Stand in for the Ulysses group and the decoder patch; return what the install was asked for."""
+    installs: list[dict] = []
+    group = object()
+
+    def install(vae, group_arg, split_dim, *, dst):
+        installs.append({"vae": vae, "group": group_arg, "split_dim": split_dim, "dst": dst})
+
+    monkeypatch.setattr(module, "install_wan_spatial_shard_decode", install)
+    monkeypatch.setattr(module.LingBotWorldCausalDMDPipeline, "_vae_shard_group", lambda self: (group, world_size))
+    return installs, group
+
+
+def test_a_multi_rank_deployment_shards_the_decoder_across_the_ulysses_ranks(monkeypatch) -> None:
+    """Multi-rank deployments shard decode by default."""
+    module = _load_pipeline_module()
+    installs, group = _record_shard_install(monkeypatch, module, world_size=2)
+
+    pipeline = module.LingBotWorldCausalDMDPipeline(od_config=_ulysses_config(2))
+
+    # Along the width, and assembled on every rank: each rank's post_decode consumes the frame.
+    assert installs == [{"vae": pipeline.vae, "group": group, "split_dim": "width", "dst": None}]
+    assert pipeline._vae_shard_split_dim == "width"
+
+
+def test_a_single_rank_deployment_leaves_the_decoder_alone(monkeypatch) -> None:
+    module = _load_pipeline_module()
+    installs, _ = _record_shard_install(monkeypatch, module, world_size=1)
+
+    pipeline = module.LingBotWorldCausalDMDPipeline(od_config=_ulysses_config(1))
+
+    assert installs == [] and pipeline._vae_shard_split_dim is None
+
+
+@pytest.mark.parametrize("enabled", [True, False])
+@pytest.mark.parametrize("width, local_width", [(832, 208), (840, 216)])
+def test_streaming_decode_reservation_uses_padded_shards(monkeypatch, enabled, width, local_width):
+    module = _load_pipeline_module()
+    installs, _ = _record_shard_install(monkeypatch, module, world_size=4)
+    config = _ulysses_config(4)
+    config.model_config.update(lingbot_vae_spatial_sharding=enabled, ar_diffusion_height=480, ar_diffusion_width=width)
+    pipeline = module.LingBotWorldCausalDMDPipeline(od_config=config)
+    from vllm_omni.experimental.ar_diffusion.streaming_decode import WanStreamingDecoder
+
+    # Exercise the real byte estimator without constructing a checkpoint VAE.
+    decoder = object.__new__(WanStreamingDecoder)
+    decoder._bytes_per_pixel_fp32 = 16.0
+    monkeypatch.setattr(pipeline, "_streaming_decoder", lambda: decoder)
+    expected_width = local_width if enabled else width
+    assert pipeline._streaming_decode_bytes_per_session() == 16 * 480 * expected_width
+    assert len(installs) == int(enabled)
+    assert pipeline._vae_shard_world_size == (4 if enabled else 1)
+    assert config.parallel_config.sequence_parallel_size == config.parallel_config.ulysses_degree == 4
+
+
+def test_disabling_vae_sharding_does_not_require_a_shard_group(monkeypatch):
+    module = _load_pipeline_module()
+    config = _ulysses_config(4)
+    config.model_config["lingbot_vae_spatial_sharding"] = False
+
+    def unexpected_group(self):
+        pytest.fail("disabled VAE sharding must not access its process group")
+
+    monkeypatch.setattr(module.LingBotWorldCausalDMDPipeline, "_vae_shard_group", unexpected_group)
+    pipeline = module.LingBotWorldCausalDMDPipeline(od_config=config)
+    assert pipeline._vae_shard_split_dim is None
+
+
+@pytest.mark.parametrize("value", ["false", 0, None])
+def test_vae_sharding_switch_requires_a_boolean(value):
+    module = _load_pipeline_module()
+    config = _ulysses_config(4)
+    config.model_config["lingbot_vae_spatial_sharding"] = value
+    with pytest.raises(ValueError, match="lingbot_vae_spatial_sharding must be a boolean"):
+        module.LingBotWorldCausalDMDPipeline(od_config=config)
+
+
+def test_vae_shard_refuses_a_group_of_the_wrong_size(monkeypatch) -> None:
+    module = _load_pipeline_module()
+    installs, _ = _record_shard_install(monkeypatch, module, world_size=4)
+
+    with pytest.raises(RuntimeError, match="sequence_parallel_size=2 but the Ulysses group has 4 ranks"):
+        module.LingBotWorldCausalDMDPipeline(od_config=_ulysses_config(2))
+    assert installs == []
+
+
 @pytest.mark.parametrize(
     ("field", "value", "feature"),
     [
@@ -758,8 +851,9 @@ def test_unsupported_parallel_modes_fail_before_component_loading(field: str, va
     assert module._loader_state.prefetch_calls == []
 
 
-def test_pure_ulysses_parallel_config_is_supported() -> None:
+def test_pure_ulysses_parallel_config_is_supported(monkeypatch) -> None:
     module = _load_pipeline_module()
+    _record_shard_install(monkeypatch, module, world_size=2)
     parallel_config = _od_config().parallel_config
     parallel_config.sequence_parallel_size = 2
     parallel_config.ulysses_degree = 2