[feature] Reduce memory usage with SP (#7728)

Signed-off-by: Maciej Bala <mbala@nvidia.com>

tests/diffusion/distributed/test_sp_storage_lifetime.py

@@ -0,0 +1,239 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import weakref
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.distributed.sp_plan import (
+    SequenceParallelConfig,
+    SequenceParallelInput,
+    SequenceParallelPartialInput,
+)
+from vllm_omni.diffusion.forward_context import get_forward_context, set_forward_context
+from vllm_omni.diffusion.hooks.base import HookRegistry
+from vllm_omni.diffusion.hooks.sequence_parallel import SequenceParallelSplitHook, apply_sequence_parallel
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion, pytest.mark.parallel]
+
+
+@pytest.fixture
+def parallel_state(monkeypatch):
+    from vllm_omni.diffusion.attention import selector
+    from vllm_omni.diffusion.distributed import parallel_state, sp_sharding
+
+    class MaskBackend:
+        @staticmethod
+        def supports_attention_mask(spec):
+            return True
+
+    def configure(size=4, rank=1):
+        monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: size)
+        monkeypatch.setattr(parallel_state, "get_sequence_parallel_rank", lambda: rank)
+        monkeypatch.setattr(parallel_state, "get_ring_parallel_world_size", lambda: 1)
+        monkeypatch.setattr(sp_sharding, "get_sequence_parallel_world_size", lambda: size)
+        monkeypatch.setattr(sp_sharding, "get_sequence_parallel_rank", lambda: rank)
+        monkeypatch.setattr(selector, "get_attn_backend_for_capability", lambda **kwargs: MaskBackend())
+
+    configure()
+    return configure
+
+
+@pytest.mark.parametrize("sequence_length", [16, 17])
+@pytest.mark.parametrize("batch", [1, 2])
+@pytest.mark.parametrize("rank", [0, 3])
+@torch.inference_mode()
+def test_cosmos3_real_sp_hook_releases_full_embedding(parallel_state, sequence_length, batch, rank):
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3GenSPPrepare, Cosmos3VFMTransformer
+
+    parallel_state(rank=rank)
+    model = nn.Module()
+    model.gen_sp_prepare = Cosmos3GenSPPrepare()
+    metadata = Cosmos3VFMTransformer._sp_plan["gen_sp_prepare"]
+    assert metadata[0].clone_shard
+    assert not metadata[1].clone_shard and not metadata[2].clone_shard
+    apply_sequence_parallel(model, SequenceParallelConfig(ulysses_degree=4), {"gen_sp_prepare": metadata})
+
+    def prepare():
+        hidden = torch.arange(batch * sequence_length * 8, dtype=torch.bfloat16).reshape(batch, sequence_length, 8)
+        full_ref = weakref.ref(hidden)
+        full_pointer = hidden.untyped_storage().data_ptr()
+        cos = torch.ones(batch, sequence_length, 1, 2)
+        sin = torch.zeros_like(cos)
+        return model.gen_sp_prepare(hidden, cos, sin), full_ref, full_pointer
+
+    for _ in range(2):
+        with set_forward_context():
+            (hidden, cos, sin), full_ref, full_pointer = prepare()
+            # Check the real backing allocation, not only Tensor.data_ptr():
+            # nonzero-offset and inference-mode views can hide the shared base.
+            assert full_ref() is None
+            assert hidden.untyped_storage().data_ptr() != full_pointer
+            assert hidden.untyped_storage().nbytes() == hidden.numel() * hidden.element_size()
+            assert hidden.is_contiguous()
+            expected = torch.arange(batch * sequence_length * 8, dtype=torch.bfloat16).reshape(
+                batch, sequence_length, 8
+            )
+            padding = -sequence_length % 4
+            expected = torch.nn.functional.pad(expected, (0, 0, 0, padding)).chunk(4, dim=1)[rank]
+            torch.testing.assert_close(hidden, expected, rtol=0, atol=0)
+            assert cos.shape[1] == sin.shape[1] == hidden.shape[1]
+
+            # Once the next layer replaces the shard, the hook must not keep
+            # even the smaller initial activation alive through later layers.
+            shard_ref = weakref.ref(hidden)
+            hidden = hidden + 1
+            assert shard_ref() is None
+
+
+@pytest.mark.parametrize("clone_shard", [False, True])
+@pytest.mark.parametrize("world_size", [1, 4])
+def test_shard_storage_policy_is_opt_in_and_skips_single_rank(parallel_state, clone_shard, world_size):
+    parallel_state(size=world_size, rank=0)
+    hook = SequenceParallelSplitHook({}, SequenceParallelConfig(ulysses_degree=4))
+    tensor = torch.arange(128).reshape(1, 16, 8)
+    spec = SequenceParallelInput(split_dim=1, clone_shard=clone_shard)
+    with set_forward_context():
+        shard = hook._prepare_sp_input(tensor, spec)
+    torch.testing.assert_close(shard, tensor.chunk(world_size, dim=1)[0])
+    owns_storage = shard.untyped_storage().data_ptr() != tensor.untyped_storage().data_ptr()
+    assert owns_storage == (clone_shard and world_size > 1)
+    if world_size == 1:
+        assert shard is tensor
+
+
+def test_owned_shard_preserves_gradients(parallel_state):
+    parallel_state(rank=1)
+    hook = SequenceParallelSplitHook({}, SequenceParallelConfig(ulysses_degree=4))
+    tensor = torch.randn(2, 16, 8, requires_grad=True)
+    with set_forward_context():
+        shard = hook._prepare_sp_input(tensor, SequenceParallelInput(split_dim=1, clone_shard=True))
+        assert shard.untyped_storage().data_ptr() != tensor.untyped_storage().data_ptr()
+        assert shard.untyped_storage().nbytes() == shard.numel() * shard.element_size()
+        shard.sum().backward()
+    expected = torch.zeros_like(tensor)
+    expected[:, 4:8] = 1
+    torch.testing.assert_close(tensor.grad, expected, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("sequence_length", [16, 17])
+@pytest.mark.parametrize("batch", [1, 2])
+@pytest.mark.parametrize("keyword", [False, True])
+@pytest.mark.parametrize("shard_group", [None, "video"])
+@torch.inference_mode()
+def test_pre_forward_clone_owns_storage_and_preserves_shard_metadata(
+    parallel_state, sequence_length, batch, keyword, shard_group
+):
+    parallel_state(rank=1)
+    full_refs = []
+    full_pointers = []
+
+    class Module(nn.Module):
+        def forward(self, x):
+            assert x.untyped_storage().data_ptr() != full_pointers[-1]
+            assert x.untyped_storage().nbytes() == x.numel() * x.element_size()
+            assert x.is_contiguous()
+            return x
+
+    module = Module()
+    spec = SequenceParallelInput(
+        split_dim=1,
+        split_output=False,
+        auto_pad=True,
+        clone_shard=True,
+        shard_group=shard_group,
+    )
+    apply_sequence_parallel(module, SequenceParallelConfig(ulysses_degree=4), {"": {"x": spec}})
+
+    def run():
+        tensor = torch.arange(batch * sequence_length * 8, dtype=torch.bfloat16).reshape(batch, sequence_length, 8)
+        full_refs.append(weakref.ref(tensor))
+        full_pointers.append(tensor.untyped_storage().data_ptr())
+        return module(x=tensor) if keyword else module(tensor)
+
+    for _ in range(2):
+        with set_forward_context():
+            ctx = get_forward_context()
+            if shard_group is not None:
+                ctx.sp_shard_metadata[shard_group] = 999  # Stale metadata must be replaced.
+            shard = run()
+            assert all(ref() is None for ref in full_refs)
+            expected = torch.arange(batch * sequence_length * 8, dtype=torch.bfloat16).reshape(
+                batch, sequence_length, 8
+            )
+            padding = -sequence_length % 4
+            expected = torch.nn.functional.pad(expected, (0, 0, 0, padding)).chunk(4, dim=1)[1]
+            torch.testing.assert_close(shard, expected, rtol=0, atol=0)
+            assert ctx.sp_shard_metadata == ({} if shard_group is None else {shard_group: sequence_length})
+            assert ctx.sp_original_seq_len == (sequence_length if padding and shard_group is None else None)
+
+
+@pytest.mark.parametrize("phase", ["success", "non_tensor", "forward_error", "split_error"])
+@pytest.mark.parametrize("keyword", [False, True])
+@torch.inference_mode()
+def test_split_hook_does_not_retain_arguments_on_any_exit(parallel_state, phase, keyword):
+    class Module(nn.Module):
+        def forward(self, x):
+            if phase == "forward_error":
+                raise RuntimeError("forward failed")
+            return None if phase == "non_tensor" else x
+
+    module = Module()
+    index = 1 if phase == "split_error" else 0
+    hook = SequenceParallelSplitHook(
+        {index: SequenceParallelInput(split_dim=1, split_output=True, clone_shard=True)},
+        SequenceParallelConfig(ulysses_degree=4),
+    )
+    HookRegistry.get_or_create(module).register_hook("sp", hook)
+    refs = []
+
+    def run():
+        tensor = torch.randn(1, 16, 8)
+        refs.append(weakref.ref(tensor))
+        return module(x=tensor) if keyword else module(tensor)
+
+    with set_forward_context():
+        if phase.endswith("error"):
+            with pytest.raises((RuntimeError, ValueError)):
+                run()
+        else:
+            result = run()
+            assert (result is None) == (phase == "non_tensor")
+    assert all(ref() is None for ref in refs)
+
+
+@pytest.mark.parametrize("keyword", [False, True])
+@pytest.mark.parametrize("tensor_source", [False, True])
+@torch.inference_mode()
+def test_partial_output_split_keeps_only_lengths_across_calls(parallel_state, keyword, tensor_source):
+    parallel_state(rank=1)
+
+    class Module(nn.Module):
+        def forward(self, x, text):
+            return x
+
+    module = Module()
+    hook = SequenceParallelSplitHook(
+        {0: SequenceParallelPartialInput(split_dim=0, text_len_source="text", split_output=True)},
+        SequenceParallelConfig(ulysses_degree=4),
+    )
+    HookRegistry.get_or_create(module).register_hook("sp", hook)
+    refs = []
+
+    def run(text_length):
+        tensor = torch.arange(text_length + 8).reshape(-1, 1)
+        text = torch.zeros(text_length, 8) if tensor_source else text_length
+        refs.append(weakref.ref(tensor))
+        if tensor_source:
+            refs.append(weakref.ref(text))
+        return module(x=tensor, text=text) if keyword else module(tensor, text)
+
+    for text_length in (2, 4):
+        with set_forward_context():
+            result = run(text_length)
+        expected = torch.cat([torch.arange(text_length), torch.arange(text_length + 2, text_length + 4)]).reshape(-1, 1)
+        torch.testing.assert_close(result, expected)
+        assert all(ref() is None for ref in refs)
+        assert hook._text_len_cache == {"text": text_length}

tests/diffusion/models/cosmos3/test_cosmos3_transformer.py

@@ -3,6 +3,7 @@
 
 from __future__ import annotations
 
+import weakref
 from types import SimpleNamespace
 from typing import Any
 
@@ -824,6 +825,134 @@ def test_cache_execution_residual_spans_final_gen_norm(monkeypatch: pytest.Monke
         assert not hasattr(model, name)
 
 
+@pytest.mark.parametrize("sequence_length", [16, 17])
+@pytest.mark.parametrize("batch", [1, 2])
+@pytest.mark.parametrize("rank", [0, 3])
+@pytest.mark.parametrize("seacache", [False, True])
+@pytest.mark.parametrize("control_count", [0, 1, 2])
+@pytest.mark.parallel
+@torch.inference_mode()
+def test_sp_releases_full_gen_embedding_during_stack(
+    monkeypatch, sequence_length, batch, rank, seacache, control_count
+) -> None:
+    from vllm_omni.diffusion.attention import selector
+    from vllm_omni.diffusion.cache.seacache import SeaCacheConfig, apply_sea_cache_hook
+    from vllm_omni.diffusion.distributed import parallel_state, sp_sharding
+    from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig
+    from vllm_omni.diffusion.forward_context import get_forward_context, set_forward_context
+    from vllm_omni.diffusion.hooks import sequence_parallel
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
+
+    full_refs = []
+    full_pointers = []
+    layer_calls = 0
+    total_length = sequence_length * (control_count + 1)
+    uses_sp = control_count <= 1
+
+    class CheckingLayer(nn.Module):
+        def forward(self, hidden, **kwargs):
+            nonlocal layer_calls
+            layer_calls += 1
+            if full_refs:
+                # Inspect lifetime inside the production stack while the
+                # caller's prepared state is still live.
+                assert full_refs[-1]() is None
+                if layer_calls % 2 == 1:
+                    assert hidden.untyped_storage().data_ptr() != full_pointers[-1]
+                assert hidden.shape == (batch, (total_length + 3) // 4, 8)
+                assert hidden.untyped_storage().nbytes() == hidden.numel() * hidden.element_size()
+                assert kwargs["freqs_cos"].shape[1] == kwargs["freqs_sin"].shape[1] == hidden.shape[1]
+            elif not uses_sp:
+                assert hidden.shape == (batch, total_length, 8)
+            return hidden + 1
+
+    monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
+    model = transformer_cosmos3.Cosmos3VFMTransformer(
+        SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
+    )
+    model.gen_layers = nn.ModuleList([CheckingLayer(), CheckingLayer()])
+    model.cached_kv = [(torch.empty(0), torch.empty(0)) for _ in model.gen_layers]
+    inputs = {
+        "hidden_states": torch.randn(batch, 2, 1, 1, sequence_length),
+        "timestep": torch.ones(batch),
+        "text_ids": torch.ones(batch, 2, dtype=torch.long),
+        "text_mask": torch.ones(batch, 2, dtype=torch.long),
+        "video_shape": (1, 1, sequence_length),
+        "control_latents": [torch.randn(batch, 2, 1, 1, sequence_length) for _ in range(control_count)],
+    }
+    gathered_reference = []
+    normalized_reference = []
+    handle = model.gen_sp_gather.register_forward_pre_hook(
+        lambda module, args: gathered_reference.append(args[0].clone())
+    )
+    norm_handle = model.norm_moe_gen.register_forward_hook(
+        lambda module, args, output: normalized_reference.append(output.clone())
+    )
+    expected = model(**inputs)
+    handle.remove()
+    norm_handle.remove()
+
+    class MaskBackend:
+        @staticmethod
+        def supports_attention_mask(spec):
+            return True
+
+    monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (4, rank, None))
+    monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 4)
+    monkeypatch.setattr(parallel_state, "get_sequence_parallel_rank", lambda: rank)
+    monkeypatch.setattr(parallel_state, "get_ring_parallel_world_size", lambda: 1)
+    monkeypatch.setattr(sp_sharding, "get_sequence_parallel_world_size", lambda: 4)
+    monkeypatch.setattr(sp_sharding, "get_sequence_parallel_rank", lambda: rank)
+    monkeypatch.setattr(selector, "get_attn_backend_for_capability", lambda **kwargs: MaskBackend())
+
+    def gather(hidden, dim, validate):
+        shard_length = hidden.shape[1]
+        start = rank * shard_length
+        valid_length = min(shard_length, total_length - start)
+        reference = normalized_reference[0] if seacache else gathered_reference[0]
+        torch.testing.assert_close(hidden[:, :valid_length], reference[:, start : start + valid_length])
+        # Supply the other ranks' outputs without retaining the original input.
+        return torch.nn.functional.pad(reference, (0, 0, 0, -total_length % 4))
+
+    monkeypatch.setattr(sequence_parallel, "sp_gather", gather)
+    sequence_parallel.apply_sequence_parallel(model, SequenceParallelConfig(ulysses_degree=4), model._sp_plan)
+
+    def record_full_embedding(module, args):
+        full_refs.append(weakref.ref(args[0]))
+        full_pointers.append(args[0].untyped_storage().data_ptr())
+
+    model.gen_sp_prepare.register_forward_pre_hook(record_full_embedding)
+    metadata = SimpleNamespace(step=0)
+    if seacache:
+        sea_hook = apply_sea_cache_hook(
+            model,
+            SeaCacheConfig(threshold=100.0, residual_order=0),
+            current_step_callback=lambda: metadata.step,
+            current_sigma_callback=lambda: 1.0,
+            num_inference_steps_callback=lambda: 4,
+        )
+        sea_hook.state_manager.set_context("cond")
+    layer_calls = 0
+    for step in range(4 if seacache else 2):
+        metadata.step = step
+        with set_forward_context():
+            output = model(**inputs)
+            torch.testing.assert_close(output, expected)
+            assert get_forward_context()._sp_shard_depth == 0
+            if seacache:
+                if uses_sp:
+                    assert full_refs[-1]() is None
+                history = sea_hook.state_manager.get_state().history
+                assert history
+                residual = history[-1][1]
+                assert residual.shape == (batch, (total_length + 3) // 4 if uses_sp else total_length, 8)
+                assert residual.untyped_storage().nbytes() == residual.numel() * residual.element_size()
+    assert layer_calls == 4
+    assert len(full_refs) == ((4 if seacache else 2) if uses_sp else 0)
+    if seacache:
+        assert sea_hook.full_count == sea_hook.skip_count == 2
+
+
 def test_no_cache_still_runs_final_gen_norm_once(monkeypatch: pytest.MonkeyPatch) -> None:
     from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
 
@@ -1394,3 +1523,28 @@ def test_compute_rope_freqs_places_text_video_action_and_sound_positions() -> No
     )
     _, offset_gen_pos = rotary.position_ids
     assert offset_gen_pos[0, 0].tolist() == [102, 103, 104, 105, 106, 107]
+
+
+def test_shard_gen_prep_rejects_inputs_already_in_execution_layout(monkeypatch: pytest.MonkeyPatch) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
+
+    monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
+    model = transformer_cosmos3.Cosmos3VFMTransformer(
+        SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
+    )
+    prep = model._gen_preprocess(
+        torch.zeros(1, 2, 1, 2, 2),
+        torch.tensor([1.0]),
+        torch.tensor([[1, 2]], dtype=torch.long),
+        torch.ones(1, 2, dtype=torch.long),
+        (1, 2, 2),
+    )
+
+    sharded = model._shard_gen_prep(prep, defer_gather=True)
+    assert sharded.freqs_gen is not None
+    assert sharded.defer_gen_gather
+    with pytest.raises(RuntimeError, match="already in the execution layout"):
+        model._shard_gen_prep(sharded)
+    # A deferred gather without sharding would gather a full-layout tensor.
+    with pytest.raises(RuntimeError, match="_shard_gen_prep"):
+        model._run_gen_stack(prep._replace(defer_gen_gather=True))