diff --git a/docs/design/feature/sequence_parallel.md b/docs/design/feature/sequence_parallel.md index 7ae263b7d..5827a2351 100644 --- a/docs/design/feature/sequence_parallel.md +++ b/docs/design/feature/sequence_parallel.md @@ -281,6 +281,10 @@ NOTE: be careful to test adequately when refactoring classes that take this styl | `expected_dims` | int \| None | Expected tensor rank for validation (optional) | | `split_output` | bool | `False`: shard **input** params; `True`: shard **output** tensors | | `auto_pad` | bool | Auto-pad if sequence not divisible by world_size (default: `False`) | +| `clone_shard` | bool | Give multi-rank shards their own contiguous storage (default: `False`); single-rank inputs are unchanged | + +With `clone_shard=True`, memory savings require callers to drop their references to the full tensor. +Other references to the same storage also keep the full allocation alive. **SequenceParallelOutput Parameters:** diff --git a/tests/diffusion/distributed/test_sp_storage_lifetime.py b/tests/diffusion/distributed/test_sp_storage_lifetime.py new file mode 100644 index 000000000..5a97556e5 --- /dev/null +++ b/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} diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py index 56c33b9e0..86898380b 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py +++ b/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)) diff --git a/vllm_omni/diffusion/cache/teacache/extractors.py b/vllm_omni/diffusion/cache/teacache/extractors.py index 2ca3e22bd..b46debb6a 100644 --- a/vllm_omni/diffusion/cache/teacache/extractors.py +++ b/vllm_omni/diffusion/cache/teacache/extractors.py @@ -1620,6 +1620,9 @@ def extract_cosmos3_context( control_weights=control_weights, transfer_share_vision_temporal_positions=transfer_share_vision_temporal_positions, ) + # Cache residuals stay rank-local: shard before the context and its + # closures retain prep; postprocess gathers on both misses and hits. + prep = module._shard_gen_prep(prep, defer_gather=True) def run_transformer_blocks() -> tuple[torch.Tensor, ...]: return (module._run_gen_stack(prep),) diff --git a/vllm_omni/diffusion/distributed/sp_plan.py b/vllm_omni/diffusion/distributed/sp_plan.py index 3e5568706..e4f97e852 100644 --- a/vllm_omni/diffusion/distributed/sp_plan.py +++ b/vllm_omni/diffusion/distributed/sp_plan.py @@ -225,6 +225,11 @@ class SequenceParallelInput: ForwardContext so models can construct attention masks when needed. Note: Ring attention does not support attention mask, so auto_pad should only be used with Ulysses SP. + clone_shard: If True, give a sharded tensor its own contiguous storage + so it does not retain the full input allocation. Applies to both + padded and unpadded splits; a single-rank input is returned unchanged. + Memory savings require the caller to drop its reference to the full + tensor; any other references to its storage also keep it allocated. shard_group: Optional key shared by tensors representing the same global sequence. Keyed groups track independent padding metadata; omitting it preserves the legacy single-sequence padding behavior. @@ -244,13 +249,15 @@ class SequenceParallelInput: expected_dims: int | None = None split_output: bool = False auto_pad: bool = False + clone_shard: bool = False shard_group: str | None = None def __repr__(self) -> str: return ( f"SequenceParallelInput(split_dim={self.split_dim}, " f"expected_dims={self.expected_dims}, split_output={self.split_output}, " - f"auto_pad={self.auto_pad}, shard_group={self.shard_group!r})" + f"auto_pad={self.auto_pad}, clone_shard={self.clone_shard}, " + f"shard_group={self.shard_group!r})" ) diff --git a/vllm_omni/diffusion/hooks/sequence_parallel.py b/vllm_omni/diffusion/hooks/sequence_parallel.py index 32f15de5f..ae4e46593 100644 --- a/vllm_omni/diffusion/hooks/sequence_parallel.py +++ b/vllm_omni/diffusion/hooks/sequence_parallel.py @@ -260,9 +260,12 @@ class SequenceParallelSplitHook(ModelHook): else: raise ValueError(f"Failed to update parameter '{name}' after sharding.") - # Store kwargs for post_forward to resolve text lengths - self._last_kwargs = kwargs - self._last_args = tuple(args_list) + # Output splitting only needs integer text lengths. Saving args/kwargs + # here would keep full hidden-state inputs alive after output sharding, + # including when the module raises before post_forward can run. + for spm in self.metadata.values(): + if isinstance(spm, SequenceParallelPartialInput) and spm.split_output: + self._resolve_text_len(spm, args_list, kwargs) return tuple(args_list), kwargs @@ -291,7 +294,7 @@ class SequenceParallelSplitHook(ModelHook): raise ValueError(f"Index {index} out of bounds for output of length {len(output_list)}.") original = output_list[index] - output_list[index] = self._prepare_sp_input(original, spm, self._last_args, self._last_kwargs) + output_list[index] = self._prepare_sp_input(original, spm) if output_list[index] is not original: actually_sharded = True equal_rank_seq_lens &= isinstance(spm, SequenceParallelInput) and spm.auto_pad @@ -392,13 +395,22 @@ class SequenceParallelSplitHook(ModelHook): if isinstance(sp_input, SequenceParallelInput): # Full split with optional auto-padding if sp_input.auto_pad: - return self._shard_with_auto_pad( + shard = self._shard_with_auto_pad( x, sp_input.split_dim, sp_input.shard_group, ) - _maybe_validate_strict_divisibility(dim=sp_input.split_dim, seq_len=x.size(sp_input.split_dim)) - return sp_shard(x, sp_input.split_dim, validate=False) + else: + _maybe_validate_strict_divisibility(dim=sp_input.split_dim, seq_len=x.size(sp_input.split_dim)) + shard = sp_shard(x, sp_input.split_dim, validate=False) + # contiguous() may be a no-op for a single-batch sequence slice; + # clone explicitly to release its full (possibly padded) backing. + if sp_input.clone_shard: + from vllm_omni.diffusion.distributed.parallel_state import get_sequence_parallel_world_size + + if get_sequence_parallel_world_size() > 1: + return shard.clone(memory_format=torch.contiguous_format) + return shard elif isinstance(sp_input, SequenceParallelPartialInput): # Partial split: keep text portion, split image portion text_len = self._resolve_text_len(sp_input, args, kwargs) diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py index b060ea420..86e3212f2 100644 --- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py @@ -1182,6 +1182,9 @@ class _GenPrepared(NamedTuple): use_multi_control_attention: bool multi_control_token_sizes: tuple[int, ...] | None multi_control_weights: tuple[float, ...] | None + # Execution-layout state; set only by Cosmos3VFMTransformer._shard_gen_prep. + freqs_gen: tuple[torch.Tensor, torch.Tensor] | None = None + defer_gen_gather: bool = False class Cosmos3VFMTransformer(nn.Module): @@ -1242,7 +1245,13 @@ class Cosmos3VFMTransformer(nn.Module): _sp_plan = { "gen_sp_prepare": { - 0: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True, auto_pad=True), + 0: SequenceParallelInput( + split_dim=1, + expected_dims=3, + split_output=True, + auto_pad=True, + clone_shard=True, + ), 1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True, auto_pad=True), 2: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True, auto_pad=True), }, @@ -1769,6 +1778,8 @@ class Cosmos3VFMTransformer(nn.Module): self, hidden_gen: torch.Tensor, *, + freqs_gen: tuple[torch.Tensor, torch.Tensor], + gather_output: bool, s_video: int, s_control: int, s_action: int, @@ -1781,13 +1792,10 @@ class Cosmos3VFMTransformer(nn.Module): multi_control_token_sizes: tuple[int, ...] | None, multi_control_weights: tuple[float, ...] | None, ) -> torch.Tensor: - """Run the complete GEN decoder between full-layout boundaries.""" - if self.cached_kv is None or self.cached_freqs_gen is None: + """Run the complete GEN decoder on inputs already in the execution layout.""" + if self.cached_kv is None: raise RuntimeError("Cosmos3 GEN cache was not initialized before running GEN layers.") - freqs_cos, freqs_sin = self.cached_freqs_gen - if not use_multi_control_attention: - hidden_gen, freqs_cos, freqs_sin = self.gen_sp_prepare(hidden_gen, freqs_cos, freqs_sin) - freqs_gen = (freqs_cos, freqs_sin) + freqs_cos, freqs_sin = freqs_gen if len(self.gen_layers) == len(self.cached_kv): for layer, (k_und, v_und) in zip(self.gen_layers, self.cached_kv, strict=True): @@ -1816,7 +1824,7 @@ class Cosmos3VFMTransformer(nn.Module): if isinstance(hidden_gen, tuple): hidden_gen = hidden_gen[0] - if not use_multi_control_attention: + if gather_output and not use_multi_control_attention: hidden_gen = self.gen_sp_gather(hidden_gen) return hidden_gen @@ -1863,8 +1871,33 @@ class Cosmos3VFMTransformer(nn.Module): control_weights=control_weights, transfer_share_vision_temporal_positions=transfer_share_vision_temporal_positions, ) + prep = self._shard_gen_prep(prep) return self._gen_postprocess(self._run_gen_stack(prep), prep) + def _shard_gen_prep(self, prep: _GenPrepared, *, defer_gather: bool = False) -> _GenPrepared: + """Move prepared GEN inputs into the execution layout. + + Callers must rebind their ``prep`` to the result before running the + stack: keeping the original would pin the full embedding even though + the sharded ``hidden_gen`` owns separate storage. + + ``defer_gather`` keeps the stack output rank-local and moves the SP + gather into ``_gen_postprocess``, so cache residuals stay sharded. + """ + if prep.freqs_gen is not None or prep.defer_gen_gather: + raise RuntimeError("Cosmos3 GEN inputs are already in the execution layout.") + if self.cached_freqs_gen is None: + raise RuntimeError("Cosmos3 GEN cache was not initialized before running GEN layers.") + if prep.use_multi_control_attention: + # Multi-control attention runs unsharded, so there is nothing to gather. + return prep._replace(freqs_gen=self.cached_freqs_gen) + hidden_gen, freqs_cos, freqs_sin = self.gen_sp_prepare(prep.hidden_gen, *self.cached_freqs_gen) + return prep._replace( + hidden_gen=hidden_gen, + freqs_gen=(freqs_cos, freqs_sin), + defer_gen_gather=defer_gather, + ) + def _gen_preprocess( self, hidden_states: torch.Tensor, @@ -2132,9 +2165,13 @@ class Cosmos3VFMTransformer(nn.Module): ) def _run_gen_stack(self, prep: _GenPrepared) -> torch.Tensor: - """Execute the cacheable full-layout GEN stack, including final norm.""" + """Execute the cacheable GEN stack, including final norm.""" + if prep.freqs_gen is None: + raise RuntimeError("Cosmos3 GEN inputs must go through _shard_gen_prep before the stack.") hidden_gen = self._run_gen_layers( prep.hidden_gen, + freqs_gen=prep.freqs_gen, + gather_output=not prep.defer_gen_gather, s_video=prep.s_video, s_control=prep.s_control, s_action=prep.s_action, @@ -2155,6 +2192,8 @@ class Cosmos3VFMTransformer(nn.Module): prep: _GenPrepared, ) -> torch.Tensor | tuple[torch.Tensor, ...]: """Project an already-normalized packed GEN state to model outputs.""" + if prep.defer_gen_gather: + hidden_gen = self.gen_sp_gather(hidden_gen) if not prep.has_action and not prep.has_sound and not prep.has_control: return self.unpatchify(self.proj_out(hidden_gen), prep.t, prep.h, prep.w)