diff --git a/docs/user_guide/diffusion/parallelism_acceleration.md b/docs/user_guide/diffusion/parallelism_acceleration.md index 324301158..007976f1a 100644 --- a/docs/user_guide/diffusion/parallelism_acceleration.md +++ b/docs/user_guide/diffusion/parallelism_acceleration.md @@ -27,7 +27,7 @@ The following table shows which models are currently supported by parallelism me | **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ❌ | | **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ❌ | | **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ❌ | -| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ❌ | ❌ | ❌ | ✅ (TP=2 only) | +| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ✅ (TP=2 only) | | **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ❌ | @@ -220,100 +220,139 @@ To measure the parallelism methods, we run benchmarks with **Qwen/Qwen-Image** m ##### How to parallelize a new model -If a diffusion model has been deployed in vLLM-Omni and supports single-card inference, you can refer to the following instructions to parallelize it with [Ulysses-SP](https://arxiv.org/pdf/2309.14509). +NOTE: "Terminology: SP vs CP" + Our "Sequence Parallelism" (SP) corresponds to "Context Parallelism" (CP) in the [diffusers library](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/_modeling_parallel.py). + We use "Sequence Parallelism" to align with vLLM-Omni's terminology. +--- -This section uses **Qwen-Image** (`QwenImageTransformer2DModel`) as the reference implementation. Qwen-Image is a **dual-stream** transformer (text + image) that performs **joint attention** across the concatenated sequences. Because of that, when enabling sequence parallel you typically: +###### Non-intrusive `_sp_plan` (Recommended) -- Chunk **image tokens** (`hidden_states`) across SP ranks along the **sequence dimension**. -- Keep **text embeddings** (`encoder_hidden_states`) **replicated** on all SP ranks for correctness (unless you implement full joint-SP semantics and explicitly split text too). -- Chunk **image RoPE freqs** to match the chunked image tokens. -- `all_gather` the final output back along the sequence dimension. +The `_sp_plan` mechanism allows SP without modifying `forward()` logic. The framework automatically registers hooks to shard inputs and gather outputs at module boundaries. -First, add the sequence-parallel helpers and (for Qwen-Image) the forward-context flag. Then, in the transformer's `forward()`: +**Requirements for `forward()` function:** -- Chunk `hidden_states` by SP world size. -- Set `get_forward_context().split_text_embed_in_sp = False` because Qwen-Image uses joint attention and we would use a full text embedding in each rank. -- After `pos_embed`, chunk `img_freqs` on `dim=0` (token axis) to match the chunked image tokens; only chunk `txt_freqs` if you explicitly split text embeddings in SP. +- Tensor operations that need sharding/gathering must happen at **`nn.Module` boundaries** (not inline Python operations) +- If your `forward()` contains inline tensor operations (e.g., `torch.cat`, `pad_sequence`) that need sharding, **extract them into a submodule** -Taking `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` as an example: +**When to create a submodule:** -```diff -from vllm_omni.diffusion.distributed.parallel_state import ( - get_sequence_parallel_rank, - get_sequence_parallel_world_size, - get_sp_group, +```python +# ❌ BAD: Inline operations - hooks cannot intercept +def forward(self, x, cap_feats): + unified = torch.cat([x, cap_feats], dim=1) # Cannot be sharded via _sp_plan + ... + +# ✅ GOOD: Extract into a submodule +class UnifiedPrepare(nn.Module): + def forward(self, x, cap_feats): + return torch.cat([x, cap_feats], dim=1) # Now can be sharded via _sp_plan + +class MyModel(nn.Module): + def __init__(self): + self.unified_prepare = UnifiedPrepare() # Submodule + + def forward(self, x, cap_feats): + unified = self.unified_prepare(x, cap_feats) # Hook can intercept here +``` + +--- + +###### Defining `_sp_plan` + +**Type definitions** (see [diffusers `_modeling_parallel.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/_modeling_parallel.py) for reference): + +```python +from vllm_omni.diffusion.distributed.sp_plan import ( + SequenceParallelInput, # Corresponds to diffusers' ContextParallelInput + SequenceParallelOutput, # Corresponds to diffusers' ContextParallelOutput ) -from vllm_omni.diffusion.forward_context import get_forward_context - -class QwenImageTransformer2DModel(...): - def forward( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor = None, - encoder_hidden_states_mask: torch.Tensor = None, - timestep: torch.LongTensor = None, - img_shapes: list[tuple[int, int, int]] | None = None, - txt_seq_lens: list[int] | None = None, - ... - ): -+ if self.parallel_config.sequence_parallel_size > 1: -+ # Chunk image tokens along sequence dimension. -+ hidden_states = torch.chunk( -+ hidden_states, -+ get_sequence_parallel_world_size(), -+ dim=-2, -+ )[get_sequence_parallel_rank()] -+ -+ # Qwen-Image uses *dual-stream* (text + image) and runs *joint attention*. -+ # Text embeddings should be replicated across SP ranks for correctness. -+ get_forward_context().split_text_embed_in_sp = False - - hidden_states = self.img_in(hidden_states) - - ... - image_rotary_emb = self.pos_embed(img_shapes, txt_seq_lens, device=hidden_states.device) - -+ def get_rotary_emb_chunk(freqs): -+ return torch.chunk(freqs, get_sequence_parallel_world_size(), dim=0)[ -+ get_sequence_parallel_rank() -+ ] -+ -+ if self.parallel_config.sequence_parallel_size > 1: -+ img_freqs, txt_freqs = image_rotary_emb -+ img_freqs = get_rotary_emb_chunk(img_freqs) -+ if get_forward_context().split_text_embed_in_sp: -+ txt_freqs = get_rotary_emb_chunk(txt_freqs) -+ image_rotary_emb = (img_freqs, txt_freqs) ``` -Next, at the end of the `forward()` function, call `get_sp_group().all_gather` to gather the chunked outputs across devices and concatenate them along the sequence dimension (matching the earlier `dim=-2` chunking): +| Parameter | Description | +|-----------|-------------| +| `split_dim` | Dimension to split/gather (usually `1` for sequence) | +| `expected_dims` | Expected tensor rank for validation (optional) | +| `split_output` | `False`: shard **input** parameters; `True`: shard **output** tensors | +| `auto_pad` | Auto-pad if sequence not divisible by world_size (Ulysses only) | -```diff -class QwenImageTransformer2DModel(...): - def forward(...): - ... - hidden_states = self.norm_out(hidden_states, temb) - output = self.proj_out(hidden_states) +**Key naming convention:** -+ if self.parallel_config.sequence_parallel_size > 1: -+ output = get_sp_group().all_gather(output, dim=-2) - return Transformer2DModelOutput(sample=output) +| Key | Meaning | Python equivalent | +|-----|---------|-------------------| +| `""` | Root model | `model` | +| `"blocks.0"` | First element of ModuleList | `model.blocks[0]` | +| `"blocks.*"` | All elements of ModuleList | `for b in model.blocks` | +| `"outputs.main"` | ModuleDict entry | `model.outputs["main"]` | + +**Dictionary key types:** + +| Key type | `split_output` | Description | +|----------|----------------|-------------| +| `"param_name"` (str) | `False` | Shard **input parameter** by name | +| `0`, `1` (int) | `True` | Shard **output tuple** by index | + +**Example** (similar to [diffusers `transformer_wan.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_wan.py)): + +```python +class MyTransformer(nn.Module): + _sp_plan = { + # Shard rope module OUTPUTS (returns tuple) + "rope": { + 0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True), # cos + 1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True), # sin + }, + # Shard transformer block INPUT parameter + "blocks.0": { + "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3), + }, + # Gather at final projection + "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), + } ``` -Finally, you can set the parallel configuration and pass it to `Omni` and start parallel inference with: -```diff -from vllm_omni import Omni -+from vllm_omni.diffusion.data import DiffusionParallelConfig -ulysses_degree = 2 +--- -omni = Omni( - model="Qwen/Qwen-Image", -+ parallel_config=DiffusionParallelConfig(ulysses_degree=2) -) +###### Hook flow -outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50) ``` +Input → [SequenceParallelSplitHook: pre_forward] → Module.forward() → [post_forward] → ... + ↓ +... → [SequenceParallelGatherHook: post_forward] → Output +``` + +1. **SplitHook** shards tensors before/after the target module +2. **Attention layers** handle Ulysses/Ring communication internally +3. **GatherHook** collects sharded outputs + +The framework automatically applies these hooks when `sequence_parallel_size > 1`. + +--- + +###### Method 2: Intrusive modification (For complex cases) + +For models with dynamic sharding logic that cannot be expressed via `_sp_plan`: + +```python +from vllm_omni.diffusion.distributed.sp_sharding import sp_shard, sp_gather + +def forward(self, hidden_states, ...): + if self.parallel_config.sequence_parallel_size > 1: + hidden_states = sp_shard(hidden_states, dim=1) + # ... computation ... + output = sp_gather(output, dim=1) + return output +``` + +--- + +###### Choosing the right approach + +| Scenario | Approach | +|----------|----------| +| Standard transformer | `_sp_plan` | +| Inline tensor ops need sharding | Extract to submodule + `_sp_plan` | +| Dynamic/conditional sharding | Intrusive modification | ### CFG-Parallel diff --git a/tests/diffusion/attention/test_sequence_parallel.py b/tests/diffusion/attention/test_attention_sp.py similarity index 100% rename from tests/diffusion/attention/test_sequence_parallel.py rename to tests/diffusion/attention/test_attention_sp.py diff --git a/tests/diffusion/distributed/test_sp_plan_hooks.py b/tests/diffusion/distributed/test_sp_plan_hooks.py new file mode 100644 index 000000000..6883ae09f --- /dev/null +++ b/tests/diffusion/distributed/test_sp_plan_hooks.py @@ -0,0 +1,1022 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the Sequence Parallelism (SP) framework. + +These tests verify the SP plan mechanism and hooks work correctly without +requiring a distributed environment. They test: +1. _sp_plan validation (sp_plan.py) +2. Hook utilities and submodule resolution (sequence_parallel.py) +3. Model _sp_plan definitions +4. Tensor sharding simulation + +Note: Our "Sequence Parallelism" (SP) corresponds to "Context Parallelism" (CP) +in diffusers. We use "Sequence Parallelism" to align with vLLM-Omni terminology. +""" + +import pytest +import torch +import torch.nn as nn + +from vllm_omni.diffusion.distributed.sp_plan import ( + SequenceParallelInput, + SequenceParallelOutput, + SequenceParallelPartialInput, + get_sp_plan_from_model, + validate_sp_plan, +) + + +def is_distributed_initialized() -> bool: + """Check if distributed environment is initialized.""" + try: + from vllm_omni.diffusion.distributed.parallel_state import get_sp_group + + get_sp_group() + return True + except (AssertionError, ImportError): + return False + + +# Decorator to skip tests that require distributed environment +requires_distributed = pytest.mark.skipif( + not is_distributed_initialized(), + reason="Requires initialized distributed environment (SP group)", +) + +# Module-level markers: these tests are diffusion + parallel related +pytestmark = [ + pytest.mark.diffusion, + pytest.mark.parallel, +] + +# ============================================================================= +# Tests for sp_plan.py +# ============================================================================= + + +@pytest.mark.cpu +class TestSequenceParallelPlanValidation: + """Test _sp_plan validation logic.""" + + def test_valid_simple_plan(self): + """Test a simple valid _sp_plan.""" + plan = { + "rope": { + 0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True), + 1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True), + }, + "blocks.0": { + "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3), + }, + "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), + } + # Should not raise + validate_sp_plan(plan) + + def test_valid_partial_input_plan(self): + """Test a valid _sp_plan with SequenceParallelPartialInput.""" + plan = { + "pos_embed": { + 0: SequenceParallelPartialInput( + split_dim=0, + text_len_source="txt_ids", + expected_dims=2, + split_output=True, + ), + }, + "blocks.0": { + "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3), + }, + } + # Should not raise + validate_sp_plan(plan) + + def test_invalid_plan_type(self): + """Test that non-dict plan raises error.""" + with pytest.raises(ValueError, match="must be a dict"): + validate_sp_plan("not a dict") + + def test_invalid_module_key_type(self): + """Test that non-string module keys raise error.""" + plan = {123: {"hidden_states": SequenceParallelInput(split_dim=1)}} + with pytest.raises(ValueError, match="keys must be strings"): + validate_sp_plan(plan) + + def test_invalid_output_index_without_split_output(self): + """Test that integer keys require split_output=True.""" + plan = { + "rope": { + 0: SequenceParallelInput(split_dim=1, split_output=False), # Invalid + } + } + with pytest.raises(ValueError, match="split_output=True"): + validate_sp_plan(plan) + + +@pytest.mark.cpu +class TestGetSpPlanFromModel: + """Test get_sp_plan_from_model utility.""" + + def test_model_with_sp_plan(self): + """Test getting _sp_plan from a model that has one.""" + + class ModelWithPlan(nn.Module): + _sp_plan = { + "layer": { + "x": SequenceParallelInput(split_dim=1), + } + } + + model = ModelWithPlan() + plan = get_sp_plan_from_model(model) + assert plan is not None + assert "layer" in plan + + def test_model_without_sp_plan(self): + """Test getting _sp_plan from a model without one.""" + + class ModelWithoutPlan(nn.Module): + pass + + model = ModelWithoutPlan() + plan = get_sp_plan_from_model(model) + assert plan is None + + +@pytest.mark.cpu +class TestSequenceParallelInputTypes: + """Test SequenceParallelInput and related types.""" + + def test_sequence_parallel_input_repr(self): + """Test SequenceParallelInput repr.""" + spi = SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True) + assert "split_dim=1" in repr(spi) + assert "expected_dims=3" in repr(spi) + assert "split_output=True" in repr(spi) + + def test_sequence_parallel_output_repr(self): + """Test SequenceParallelOutput repr.""" + spo = SequenceParallelOutput(gather_dim=1, expected_dims=3) + assert "gather_dim=1" in repr(spo) + assert "expected_dims=3" in repr(spo) + + def test_sequence_parallel_partial_input_repr(self): + """Test SequenceParallelPartialInput repr.""" + sppi = SequenceParallelPartialInput( + split_dim=0, + text_len_source="txt_ids", + expected_dims=2, + split_output=True, + ) + assert "split_dim=0" in repr(sppi) + assert "txt_ids" in repr(sppi) + assert "expected_dims=2" in repr(sppi) + assert "split_output=True" in repr(sppi) + + def test_sequence_parallel_partial_input_with_int_source(self): + """Test SequenceParallelPartialInput with integer text_len_source.""" + sppi = SequenceParallelPartialInput( + split_dim=0, + text_len_source=512, # Fixed length + expected_dims=2, + ) + assert sppi.text_len_source == 512 + + +# ============================================================================= +# Tests for sequence_parallel.py +# ============================================================================= + + +@pytest.mark.cpu +class TestModuleForwardMetadata: + """Test ModuleForwardMetadata parameter resolution.""" + + def test_get_parameter_from_kwargs(self): + """Test getting parameter from kwargs.""" + from vllm_omni.diffusion.hooks.sequence_parallel import ModuleForwardMetadata + + class DummyModule(nn.Module): + def forward(self, hidden_states, encoder_hidden_states): + pass + + metadata = ModuleForwardMetadata() + metadata._cls = DummyModule + + kwargs = {"hidden_states": torch.randn(2, 4, 8)} + val, is_kwarg, index = metadata._get_parameter_from_args_kwargs("hidden_states", (), kwargs) + assert is_kwarg is True + assert index is None + assert val.shape == (2, 4, 8) + + def test_get_parameter_from_args(self): + """Test getting parameter from positional args.""" + from vllm_omni.diffusion.hooks.sequence_parallel import ModuleForwardMetadata + + class DummyModule(nn.Module): + def forward(self, hidden_states, encoder_hidden_states): + pass + + metadata = ModuleForwardMetadata() + metadata._cls = DummyModule + + tensor = torch.randn(2, 4, 8) + args = (tensor,) + val, is_kwarg, index = metadata._get_parameter_from_args_kwargs("hidden_states", args, {}) + assert is_kwarg is False + assert index == 0 + assert torch.equal(val, tensor) + + def test_parameter_caching(self): + """Test that parameter indices are cached.""" + from vllm_omni.diffusion.hooks.sequence_parallel import ModuleForwardMetadata + + class DummyModule(nn.Module): + def forward(self, a, b, c): + pass + + metadata = ModuleForwardMetadata() + metadata._cls = DummyModule + + # First call - should populate cache + args = (torch.randn(1), torch.randn(1), torch.randn(1)) + metadata._get_parameter_from_args_kwargs("b", args, {}) + + # Check cache was populated + assert metadata.cached_parameter_indices is not None + assert metadata.cached_parameter_indices["a"] == 0 + assert metadata.cached_parameter_indices["b"] == 1 + assert metadata.cached_parameter_indices["c"] == 2 + + +@pytest.mark.cpu +class TestGetSubmoduleByName: + """Test _get_submodule_by_name function.""" + + def test_root_module(self): + """Test getting root module with empty string.""" + from vllm_omni.diffusion.hooks.sequence_parallel import _get_submodule_by_name + + model = nn.Linear(10, 10) + submodule = _get_submodule_by_name(model, "") + assert submodule is model + + def test_simple_submodule(self): + """Test getting a simple submodule.""" + from vllm_omni.diffusion.hooks.sequence_parallel import _get_submodule_by_name + + class Model(nn.Module): + def __init__(self): + super().__init__() + self.layer = nn.Linear(10, 10) + + model = Model() + submodule = _get_submodule_by_name(model, "layer") + assert submodule is model.layer + + def test_nested_submodule(self): + """Test getting a nested submodule.""" + from vllm_omni.diffusion.hooks.sequence_parallel import _get_submodule_by_name + + class Model(nn.Module): + def __init__(self): + super().__init__() + self.encoder = nn.Sequential(nn.Linear(10, 10), nn.ReLU()) + + model = Model() + submodule = _get_submodule_by_name(model, "encoder.0") + assert isinstance(submodule, nn.Linear) + + def test_module_list_by_index(self): + """Test getting element from ModuleList by index.""" + from vllm_omni.diffusion.hooks.sequence_parallel import _get_submodule_by_name + + class Model(nn.Module): + def __init__(self): + super().__init__() + self.blocks = nn.ModuleList([nn.Linear(10, 10) for _ in range(3)]) + + model = Model() + submodule = _get_submodule_by_name(model, "blocks.0") + assert submodule is model.blocks[0] + + submodule = _get_submodule_by_name(model, "blocks.2") + assert submodule is model.blocks[2] + + def test_wildcard_modulelist(self): + """Test wildcard matching for ModuleList.""" + from vllm_omni.diffusion.hooks.sequence_parallel import _get_submodule_by_name + + class Model(nn.Module): + def __init__(self): + super().__init__() + self.blocks = nn.ModuleList([nn.Linear(10, 10) for _ in range(3)]) + + model = Model() + submodules = _get_submodule_by_name(model, "blocks.*") + assert isinstance(submodules, list) + assert len(submodules) == 3 + for i, sm in enumerate(submodules): + assert sm is model.blocks[i] + + def test_module_dict(self): + """Test getting submodule from ModuleDict.""" + from vllm_omni.diffusion.hooks.sequence_parallel import _get_submodule_by_name + + class Model(nn.Module): + def __init__(self): + super().__init__() + self.outputs = nn.ModuleDict({"main": nn.Linear(10, 10), "aux": nn.Linear(10, 5)}) + + model = Model() + submodule = _get_submodule_by_name(model, "outputs.main") + assert submodule is model.outputs["main"] + + submodule = _get_submodule_by_name(model, "outputs.aux") + assert submodule is model.outputs["aux"] + + def test_invalid_submodule_raises(self): + """Test that invalid submodule path raises error.""" + from vllm_omni.diffusion.hooks.sequence_parallel import _get_submodule_by_name + + class Model(nn.Module): + def __init__(self): + super().__init__() + self.layer = nn.Linear(10, 10) + + model = Model() + with pytest.raises(ValueError, match="not a submodule"): + _get_submodule_by_name(model, "nonexistent") + + def test_multiple_wildcards_raises(self): + """Test that multiple wildcards raise error.""" + from vllm_omni.diffusion.hooks.sequence_parallel import _get_submodule_by_name + + model = nn.Linear(10, 10) + with pytest.raises(ValueError, match="only be used once"): + _get_submodule_by_name(model, "a.*.b.*") + + +@pytest.mark.cpu +class TestHookRegistration: + """Test hook registration logic (without distributed backend).""" + + def test_plan_validation_before_apply(self): + """Test that invalid plans are rejected before hook registration.""" + + class SimpleModel(nn.Module): + def __init__(self): + super().__init__() + self.proj_in = nn.Linear(10, 10) + self.proj_out = nn.Linear(10, 10) + + def forward(self, x): + return self.proj_out(self.proj_in(x)) + + # Invalid plan (non-string key) + invalid_plan = { + 123: {"x": SequenceParallelInput(split_dim=1)}, + } + + with pytest.raises(ValueError): + validate_sp_plan(invalid_plan) + + def test_valid_plan_structure_for_model(self): + """Test that a valid plan can be defined for a model.""" + + class SimpleModel(nn.Module): + _sp_plan = { + "proj_in": {"x": SequenceParallelInput(split_dim=1, expected_dims=3)}, + "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), + } + + def __init__(self): + super().__init__() + self.proj_in = nn.Linear(10, 10) + self.proj_out = nn.Linear(10, 10) + + def forward(self, x): + return self.proj_out(self.proj_in(x)) + + model = SimpleModel() + plan = get_sp_plan_from_model(model) + + assert plan is not None + assert "proj_in" in plan + assert "proj_out" in plan + + # Verify submodules exist + from vllm_omni.diffusion.hooks.sequence_parallel import _get_submodule_by_name + + assert _get_submodule_by_name(model, "proj_in") is model.proj_in + assert _get_submodule_by_name(model, "proj_out") is model.proj_out + + +# ============================================================================= +# Tests for model _sp_plan definitions +# ============================================================================= + + +@pytest.mark.L4 +class TestModelSpPlans: + """Test that model _sp_plan definitions are valid. + + These tests import actual model classes to verify _sp_plan structure. + May require GPU for model imports. + """ + + def test_zimage_transformer_sp_plan(self): + """Test ZImageTransformer2DModel _sp_plan structure. + + The plan specifies: + - unified_prepare: Shard all 4 outputs (unified, cos, sin, attn_mask) + - all_final_layer.2-1: Gather outputs after final layer + + Note: _sp_plan corresponds to diffusers' _cp_plan (Context Parallelism) + """ + try: + from vllm_omni.diffusion.models.z_image.z_image_transformer import ZImageTransformer2DModel + + plan = getattr(ZImageTransformer2DModel, "_sp_plan", None) + assert plan is not None, "ZImageTransformer2DModel should define _sp_plan" + assert isinstance(plan, dict) + + assert "unified_prepare" in plan + unified_prepare_plan = plan["unified_prepare"] + # Check all 4 outputs are sharded with split_output=True + assert 0 in unified_prepare_plan # unified + assert 1 in unified_prepare_plan # unified_cos + assert 2 in unified_prepare_plan # unified_sin + assert 3 in unified_prepare_plan # unified_attn_mask + + # Check output gathering + assert "all_final_layer.2-1" in plan + + validate_sp_plan(plan) + except ImportError: + pytest.skip("ZImageTransformer2DModel not available") + + def test_qwen_image_transformer_sp_plan(self): + """Test QwenImageTransformer2DModel _sp_plan structure. + + Qwen-Image follows the diffusers pattern similar to Z-Image: + - image_rope_prepare: Shards hidden_states and vid_freqs together + - proj_out: Gathers output + + Key insight: hidden_states and vid_freqs MUST be sharded together + to maintain dimension alignment for RoPE computation. + + Note: _sp_plan corresponds to diffusers' _cp_plan (Context Parallelism) + """ + try: + from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import ( + QwenImageTransformer2DModel, + ) + + plan = getattr(QwenImageTransformer2DModel, "_sp_plan", None) + assert plan is not None, "QwenImageTransformer2DModel should define _sp_plan" + assert isinstance(plan, dict) + + # Check image_rope_prepare sharding + assert "image_rope_prepare" in plan + rope_plan = plan["image_rope_prepare"] + # hidden_states (index 0) + assert 0 in rope_plan + assert rope_plan[0].split_dim == 1 + assert rope_plan[0].split_output is True + # vid_freqs (index 1) + assert 1 in rope_plan + assert rope_plan[1].split_dim == 0 + assert rope_plan[1].split_output is True + # txt_freqs (index 2) should NOT be in plan (kept replicated) + assert 2 not in rope_plan + + # Check output gathering at proj_out + assert "proj_out" in plan + proj_out_plan = plan["proj_out"] + assert proj_out_plan.gather_dim == 1 + + validate_sp_plan(plan) + except ImportError: + pytest.skip("QwenImageTransformer2DModel not available") + + +# ============================================================================= +# Tests for tensor sharding simulation (no distributed required) +# ============================================================================= + + +@pytest.mark.cpu +class TestMockSharding: + """Test tensor sharding logic (mocked, no distributed).""" + + def test_shard_tensor_simulation(self): + """Simulate tensor sharding without distributed backend.""" + # Create a test tensor + batch_size, seq_len, hidden_dim = 2, 16, 64 + tensor = torch.randn(batch_size, seq_len, hidden_dim) + + # Simulate sharding for world_size=4 + world_size = 4 + rank = 1 + + # Manual chunking (what sp_shard does internally) + chunks = tensor.chunk(world_size, dim=1) + sharded = chunks[rank] + + assert sharded.shape == (batch_size, seq_len // world_size, hidden_dim) + assert sharded.shape == (2, 4, 64) + + def test_partial_shard_simulation(self): + """Simulate partial sharding (text kept, image sharded).""" + # Create a test tensor with [text, image] concatenated + batch_size = 2 + text_len = 8 + image_len = 16 + hidden_dim = 64 + + text_part = torch.randn(batch_size, text_len, hidden_dim) + image_part = torch.randn(batch_size, image_len, hidden_dim) + tensor = torch.cat([text_part, image_part], dim=1) + + assert tensor.shape == (batch_size, text_len + image_len, hidden_dim) + + # Simulate partial sharding for world_size=4, rank=1 + world_size = 4 + rank = 1 + dim = 1 + + # Extract parts + text_kept = tensor.narrow(dim, 0, text_len) + image_full = tensor.narrow(dim, text_len, image_len) + + # Shard only image part + image_chunks = image_full.chunk(world_size, dim=dim) + image_sharded = image_chunks[rank] + + # Concatenate back + result = torch.cat([text_kept, image_sharded], dim=dim) + + expected_len = text_len + image_len // world_size + assert result.shape == (batch_size, expected_len, hidden_dim) + assert result.shape == (2, 8 + 4, 64) # text_len + image_len/4 + + def test_gather_tensor_simulation(self): + """Simulate tensor gathering without distributed backend.""" + # Create sharded tensors (as if from different ranks) + batch_size, shard_seq_len, hidden_dim = 2, 4, 64 + world_size = 4 + + shards = [torch.randn(batch_size, shard_seq_len, hidden_dim) for _ in range(world_size)] + + # Simulate gathering (concatenate along dim 1) + gathered = torch.cat(shards, dim=1) + + assert gathered.shape == (batch_size, shard_seq_len * world_size, hidden_dim) + assert gathered.shape == (2, 16, 64) + + def test_padding_simulation(self): + """Simulate padding for non-divisible sequence lengths.""" + # Create tensor with non-divisible sequence length + batch_size, seq_len, hidden_dim = 2, 17, 64 # 17 not divisible by 4 + tensor = torch.randn(batch_size, seq_len, hidden_dim) + + world_size = 4 + dim = 1 + + # Calculate padding needed + remainder = seq_len % world_size + if remainder != 0: + pad_size = world_size - remainder + else: + pad_size = 0 + + assert pad_size == 3 # 17 + 3 = 20, divisible by 4 + + # Pad tensor + if pad_size > 0: + pad_shape = list(tensor.shape) + pad_shape[dim] = pad_size + padding = torch.zeros(pad_shape, dtype=tensor.dtype, device=tensor.device) + padded = torch.cat([tensor, padding], dim=dim) + else: + padded = tensor + + assert padded.shape == (batch_size, seq_len + pad_size, hidden_dim) + assert padded.shape == (2, 20, 64) + + # Now can shard evenly + chunks = padded.chunk(world_size, dim=dim) + assert all(c.shape == (2, 5, 64) for c in chunks) + + +# ============================================================================= +# Additional tests for sequence_parallel.py coverage +# ============================================================================= + + +@pytest.mark.cpu +class TestUnwrapModule: + """Test _unwrap_module function.""" + + def test_unwrap_simple_module(self): + """Test that a simple module returns itself.""" + from vllm_omni.diffusion.hooks.sequence_parallel import _unwrap_module + + module = nn.Linear(10, 10) + result = _unwrap_module(module) + assert result is module + + def test_unwrap_sequential_single(self): + """Test unwrapping a Sequential with single child.""" + from vllm_omni.diffusion.hooks.sequence_parallel import _unwrap_module + + inner = nn.Linear(10, 10) + wrapper = nn.Sequential(inner) + result = _unwrap_module(wrapper) + # Should unwrap to the inner module + assert result is inner + + def test_unwrap_nested_wrapper(self): + """Test unwrapping nested single-child wrappers.""" + from vllm_omni.diffusion.hooks.sequence_parallel import _unwrap_module + + inner = nn.Linear(10, 10) + wrapper1 = nn.Sequential(inner) + wrapper2 = nn.Sequential(wrapper1) + result = _unwrap_module(wrapper2) + # Should fully unwrap to the innermost module + assert result is inner + + +@pytest.mark.cpu +class TestSequenceParallelSplitHookInit: + """Test SequenceParallelSplitHook initialization and setup.""" + + def test_hook_init(self): + """Test SequenceParallelSplitHook initialization.""" + from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig + from vllm_omni.diffusion.hooks.sequence_parallel import SequenceParallelSplitHook + + metadata = { + "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3), + } + config = SequenceParallelConfig(ulysses_degree=2, ring_degree=1) + + hook = SequenceParallelSplitHook(metadata, config) + assert hook.metadata == metadata + assert hook.config == config + assert hook.module_forward_metadata is None # Not initialized until initialize_hook + + def test_hook_initialize(self): + """Test SequenceParallelSplitHook.initialize_hook.""" + from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig + from vllm_omni.diffusion.hooks.sequence_parallel import SequenceParallelSplitHook + + class DummyModule(nn.Module): + def forward(self, hidden_states, encoder_hidden_states): + return hidden_states + + metadata = { + "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3), + } + config = SequenceParallelConfig(ulysses_degree=2, ring_degree=1) + + hook = SequenceParallelSplitHook(metadata, config) + module = DummyModule() + + # Initialize hook + result = hook.initialize_hook(module) + assert result is module + assert hook.module_forward_metadata is not None + assert hook.module_forward_metadata._cls is DummyModule + + +@pytest.mark.cpu +class TestSequenceParallelGatherHookInit: + """Test SequenceParallelGatherHook initialization.""" + + def test_hook_init_single_output(self): + """Test SequenceParallelGatherHook with single output.""" + from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig + from vllm_omni.diffusion.hooks.sequence_parallel import SequenceParallelGatherHook + + metadata = SequenceParallelOutput(gather_dim=1, expected_dims=3) + config = SequenceParallelConfig(ulysses_degree=2, ring_degree=1) + + hook = SequenceParallelGatherHook(metadata, config) + # Single output should be wrapped in a list + assert isinstance(hook.metadata, list) + assert len(hook.metadata) == 1 + assert hook.metadata[0].gather_dim == 1 + + def test_hook_init_multiple_outputs(self): + """Test SequenceParallelGatherHook with multiple outputs.""" + from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig + from vllm_omni.diffusion.hooks.sequence_parallel import SequenceParallelGatherHook + + metadata = [ + SequenceParallelOutput(gather_dim=1, expected_dims=3), + SequenceParallelOutput(gather_dim=2, expected_dims=4), + ] + config = SequenceParallelConfig(ulysses_degree=2, ring_degree=1) + + hook = SequenceParallelGatherHook(metadata, config) + assert len(hook.metadata) == 2 + assert hook.metadata[0].gather_dim == 1 + assert hook.metadata[1].gather_dim == 2 + + +@pytest.mark.cpu +class TestResolveTextLen: + """Test _resolve_text_len in SequenceParallelSplitHook.""" + + def test_resolve_int_source(self): + """Test resolving text length from integer source.""" + from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig + from vllm_omni.diffusion.hooks.sequence_parallel import SequenceParallelSplitHook + + class DummyModule(nn.Module): + def forward(self, x, txt_ids): + return x + + partial_input = SequenceParallelPartialInput( + split_dim=1, + text_len_source=256, # Fixed integer + expected_dims=3, + ) + config = SequenceParallelConfig(ulysses_degree=2, ring_degree=1) + + hook = SequenceParallelSplitHook({"x": partial_input}, config) + hook.initialize_hook(DummyModule()) + + # Resolve with integer source + text_len = hook._resolve_text_len(partial_input, (), {}) + assert text_len == 256 + + def test_resolve_string_source_from_tensor(self): + """Test resolving text length from tensor parameter.""" + from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig + from vllm_omni.diffusion.hooks.sequence_parallel import SequenceParallelSplitHook + + class DummyModule(nn.Module): + def forward(self, x, txt_ids): + return x + + partial_input = SequenceParallelPartialInput( + split_dim=1, + text_len_source="txt_ids", # Get from parameter + expected_dims=3, + ) + config = SequenceParallelConfig(ulysses_degree=2, ring_degree=1) + + hook = SequenceParallelSplitHook({"x": partial_input}, config) + hook.initialize_hook(DummyModule()) + + # Provide txt_ids tensor + txt_ids = torch.randn(128, 64) # shape[0] = 128 + kwargs = {"txt_ids": txt_ids} + + text_len = hook._resolve_text_len(partial_input, (), kwargs) + assert text_len == 128 + + def test_resolve_text_len_caching(self): + """Test that text length is cached.""" + from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig + from vllm_omni.diffusion.hooks.sequence_parallel import SequenceParallelSplitHook + + class DummyModule(nn.Module): + def forward(self, x, txt_ids): + return x + + partial_input = SequenceParallelPartialInput( + split_dim=1, + text_len_source="txt_ids", + expected_dims=3, + ) + config = SequenceParallelConfig(ulysses_degree=2, ring_degree=1) + + hook = SequenceParallelSplitHook({"x": partial_input}, config) + hook.initialize_hook(DummyModule()) + + txt_ids = torch.randn(64, 32) + kwargs = {"txt_ids": txt_ids} + + # First call - should populate cache + hook._resolve_text_len(partial_input, (), kwargs) + assert "txt_ids" in hook._text_len_cache + assert hook._text_len_cache["txt_ids"] == 64 + + # Second call - should use cache + text_len = hook._resolve_text_len(partial_input, (), kwargs) + assert text_len == 64 + + +@pytest.mark.cpu +class TestHookNameTemplates: + """Test hook name template generation.""" + + def test_input_hook_name(self): + """Test input hook name format.""" + from vllm_omni.diffusion.hooks.sequence_parallel import _SP_INPUT_HOOK_TEMPLATE + + name = _SP_INPUT_HOOK_TEMPLATE.format("blocks.0") + assert name == "sp_input---blocks.0" + + def test_output_hook_name(self): + """Test output hook name format.""" + from vllm_omni.diffusion.hooks.sequence_parallel import _SP_OUTPUT_HOOK_TEMPLATE + + name = _SP_OUTPUT_HOOK_TEMPLATE.format("proj_out") + assert name == "sp_output---proj_out" + + +@pytest.mark.cpu +class TestApplyRemoveSequenceParallel: + """Test apply_sequence_parallel and remove_sequence_parallel functions.""" + + def test_apply_sp_registers_hooks(self): + """Test that apply_sequence_parallel registers hooks on modules.""" + from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig + from vllm_omni.diffusion.hooks.sequence_parallel import ( + _SP_INPUT_HOOK_TEMPLATE, + _SP_OUTPUT_HOOK_TEMPLATE, + apply_sequence_parallel, + ) + + class SimpleModel(nn.Module): + def __init__(self): + super().__init__() + self.proj_in = nn.Linear(10, 10) + self.proj_out = nn.Linear(10, 10) + + def forward(self, hidden_states): + x = self.proj_in(hidden_states) + return self.proj_out(x) + + model = SimpleModel() + config = SequenceParallelConfig(ulysses_degree=2, ring_degree=1) + plan = { + "proj_in": {"hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3)}, + "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), + } + + # Apply SP + apply_sequence_parallel(model, config, plan) + + # Check hooks are registered + + assert hasattr(model.proj_in, "_hook_registry") + assert hasattr(model.proj_out, "_hook_registry") + + proj_in_registry = model.proj_in._hook_registry + proj_out_registry = model.proj_out._hook_registry + + assert _SP_INPUT_HOOK_TEMPLATE.format("proj_in") in proj_in_registry._hooks + assert _SP_OUTPUT_HOOK_TEMPLATE.format("proj_out") in proj_out_registry._hooks + + def test_remove_sp_removes_hooks(self): + """Test that remove_sequence_parallel removes hooks from modules.""" + from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig + from vllm_omni.diffusion.hooks.sequence_parallel import ( + _SP_INPUT_HOOK_TEMPLATE, + _SP_OUTPUT_HOOK_TEMPLATE, + apply_sequence_parallel, + remove_sequence_parallel, + ) + + class SimpleModel(nn.Module): + def __init__(self): + super().__init__() + self.proj_in = nn.Linear(10, 10) + self.proj_out = nn.Linear(10, 10) + + def forward(self, hidden_states): + x = self.proj_in(hidden_states) + return self.proj_out(x) + + model = SimpleModel() + config = SequenceParallelConfig(ulysses_degree=2, ring_degree=1) + plan = { + "proj_in": {"hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3)}, + "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), + } + + # Apply then remove SP + apply_sequence_parallel(model, config, plan) + remove_sequence_parallel(model, plan) + + # Check hooks are removed + proj_in_registry = model.proj_in._hook_registry + proj_out_registry = model.proj_out._hook_registry + + assert _SP_INPUT_HOOK_TEMPLATE.format("proj_in") not in proj_in_registry._hooks + assert _SP_OUTPUT_HOOK_TEMPLATE.format("proj_out") not in proj_out_registry._hooks + + def test_apply_sp_with_wildcard(self): + """Test apply_sequence_parallel with wildcard module names.""" + from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig + from vllm_omni.diffusion.hooks.sequence_parallel import ( + _SP_INPUT_HOOK_TEMPLATE, + apply_sequence_parallel, + ) + + class Block(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(10, 10) + + def forward(self, x): + return self.linear(x) + + class Model(nn.Module): + def __init__(self): + super().__init__() + self.blocks = nn.ModuleList([Block() for _ in range(3)]) + + def forward(self, x): + for block in self.blocks: + x = block(x) + return x + + model = Model() + config = SequenceParallelConfig(ulysses_degree=2, ring_degree=1) + plan = { + "blocks.*": {"x": SequenceParallelInput(split_dim=1, expected_dims=3)}, + } + + # Apply SP + apply_sequence_parallel(model, config, plan) + + # Check all blocks have hooks registered + for i, block in enumerate(model.blocks): + assert hasattr(block, "_hook_registry") + registry = block._hook_registry + assert _SP_INPUT_HOOK_TEMPLATE.format("blocks.*") in registry._hooks + + +@pytest.mark.cpu +class TestDimensionValidation: + """Test expected_dims validation in hooks.""" + + def test_skip_shard_on_wrong_dims(self): + """Test that sharding is skipped when tensor dims don't match expected.""" + from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig + from vllm_omni.diffusion.hooks.sequence_parallel import SequenceParallelSplitHook + + class DummyModule(nn.Module): + def forward(self, x): + return x + + # Expect 3D tensor + metadata = { + "x": SequenceParallelInput(split_dim=1, expected_dims=3), + } + config = SequenceParallelConfig(ulysses_degree=2, ring_degree=1) + + hook = SequenceParallelSplitHook(metadata, config) + hook.initialize_hook(DummyModule()) + + # Provide 4D tensor (wrong dims) + tensor_4d = torch.randn(2, 4, 8, 16) + + # _prepare_sp_input should return tensor unchanged when dims don't match + result = hook._prepare_sp_input(tensor_4d, metadata["x"], (), {}) + # Since expected_dims=3 but tensor has 4 dims, should return original + assert result.shape == tensor_4d.shape + + +@pytest.mark.cpu +class TestSequenceParallelConfig: + """Test SequenceParallelConfig dataclass.""" + + def test_config_defaults_invalid(self): + """Test that SequenceParallelConfig with default values raises error. + + At least one of ulysses_degree or ring_degree must be > 1 to enable SP. + """ + from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig + + with pytest.raises(ValueError, match="must be > 1"): + SequenceParallelConfig() # Both defaults are 1, which is invalid + + def test_config_ulysses_only(self): + """Test SequenceParallelConfig with Ulysses only.""" + from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig + + config = SequenceParallelConfig(ulysses_degree=4, ring_degree=1) + assert config.sequence_parallel_size == 4 + + def test_config_ring_only(self): + """Test SequenceParallelConfig with Ring only.""" + from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig + + config = SequenceParallelConfig(ulysses_degree=1, ring_degree=4) + assert config.sequence_parallel_size == 4 + + def test_config_hybrid(self): + """Test SequenceParallelConfig with hybrid (Ulysses + Ring).""" + from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig + + config = SequenceParallelConfig(ulysses_degree=2, ring_degree=4) + assert config.sequence_parallel_size == 8 diff --git a/tests/e2e/offline_inference/test_sequence_parallel.py b/tests/e2e/offline_inference/test_sequence_parallel.py index 29172b7d1..9de9fd981 100644 --- a/tests/e2e/offline_inference/test_sequence_parallel.py +++ b/tests/e2e/offline_inference/test_sequence_parallel.py @@ -9,6 +9,7 @@ correctly with diffusion models. It uses minimal settings to keep test time short for CI. """ +import gc import os import sys import time @@ -21,7 +22,7 @@ import torch.distributed as dist from PIL import Image # ruff: noqa: E402 -REPO_ROOT = Path(__file__).resolve().parents[2] +REPO_ROOT = Path(__file__).resolve().parents[3] if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) @@ -30,7 +31,6 @@ from vllm_omni.diffusion.data import DiffusionParallelConfig from vllm_omni.diffusion.distributed.parallel_state import device_count from vllm_omni.diffusion.envs import get_device_name -# Use random weights model for testing models = ["riverclouds/qwen_image_random"] PROMPT = "a photo of a cat sitting on a laptop keyboard" @@ -51,44 +51,145 @@ def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]: return abs_diff.mean().item(), abs_diff.max().item() -def _get_images(output): - """Extract images from output, handling both dict and SimpleNamespace types. +def _cleanup_distributed(): + """Clean up distributed environment and GPU resources.""" + if dist.is_initialized(): + dist.destroy_process_group() - The output structure varies depending on serialization path: - - Direct memory: SimpleNamespace with .images attribute - - SHM serialization: dict with "images" key (dataclass converted via asdict) - - Wrapped output: SimpleNamespace(output=...) which needs unwrapping - """ - # Check if output has direct images attribute (diffusion mode) - if hasattr(output, "images") and output.images: - return output.images + for key in ["MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"]: + os.environ.pop(key, None) - # Check request_output for pipeline mode - if output.request_output is None: - return None + gc.collect() - if isinstance(output.request_output, list) and len(output.request_output) == 0: - return None + if torch.cuda.is_available(): + torch.cuda.empty_cache() + torch.cuda.synchronize() - item = output.request_output[0] + time.sleep(5) - # Handle wrapped SimpleNamespace (e.g. from omni_stage.py) - # Some items are wrapped as SimpleNamespace(request_id=..., output=...) - while hasattr(item, "output") and not hasattr(item, "images"): - item = item.output - if item is None: - return None - # Handle both dict (from SHM serialization) and object (direct) types - if isinstance(item, dict): - return item.get("images") - return getattr(item, "images", None) +def _run_baseline(model_name: str, dtype: torch.dtype, attn_backend: str, height: int, width: int, seed: int): + """Run baseline inference (no SP).""" + baseline_parallel_config = DiffusionParallelConfig(ulysses_degree=1, ring_degree=1) + baseline = Omni( + model=model_name, + parallel_config=baseline_parallel_config, + dtype=dtype, + attention_backend=attn_backend, + ) + + try: + outputs = baseline.generate( + PROMPT, + height=height, + width=width, + num_inference_steps=4, + guidance_scale=0.0, + generator=torch.Generator(get_device_name()).manual_seed(seed), + num_outputs_per_prompt=1, + ) + return outputs[0].request_output[0].images + finally: + baseline.close() + _cleanup_distributed() + + +def _run_sp( + model_name: str, + dtype: torch.dtype, + attn_backend: str, + height: int, + width: int, + seed: int, + ulysses_degree: int, + ring_degree: int, +): + """Run SP inference.""" + sp_parallel_config = DiffusionParallelConfig(ulysses_degree=ulysses_degree, ring_degree=ring_degree) + sp = Omni( + model=model_name, + parallel_config=sp_parallel_config, + dtype=dtype, + attention_backend=attn_backend, + ) + + try: + outputs = sp.generate( + PROMPT, + height=height, + width=width, + num_inference_steps=4, + guidance_scale=0.0, + generator=torch.Generator(get_device_name()).manual_seed(seed), + num_outputs_per_prompt=1, + ) + return outputs[0].request_output[0].images + finally: + sp.close() + _cleanup_distributed() + + +@pytest.mark.parametrize("model_name", models) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize("attn_backend", ["sdpa"]) +def test_baseline_only(model_name: str, dtype: torch.dtype, attn_backend: str): + """Test baseline inference only (no SP).""" + height = 256 + width = 256 + seed = 42 + + images = _run_baseline(model_name, dtype, attn_backend, height, width, seed) + + assert images is not None + assert len(images) == 1 + assert images[0].width == width + assert images[0].height == height + + +@pytest.mark.parametrize("model_name", models) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize("attn_backend", ["sdpa"]) +def test_sp_ulysses2_only(model_name: str, dtype: torch.dtype, attn_backend: str): + """Test SP inference only (ulysses=2).""" + if device_count() < 2: + pytest.skip(f"Test requires 2 GPUs but only {device_count()} available") + + height = 256 + width = 256 + seed = 42 + + images = _run_sp(model_name, dtype, attn_backend, height, width, seed, ulysses_degree=2, ring_degree=1) + + assert images is not None + assert len(images) == 1 + assert images[0].width == width + assert images[0].height == height + + +@pytest.mark.parametrize("model_name", models) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize("attn_backend", ["sdpa"]) +def test_sp_ring2_only(model_name: str, dtype: torch.dtype, attn_backend: str): + """Test SP inference only (ring=2).""" + if device_count() < 2: + pytest.skip(f"Test requires 2 GPUs but only {device_count()} available") + + height = 256 + width = 256 + seed = 42 + + images = _run_sp(model_name, dtype, attn_backend, height, width, seed, ulysses_degree=1, ring_degree=2) + + assert images is not None + assert len(images) == 1 + assert images[0].width == width + assert images[0].height == height @pytest.mark.parametrize("model_name", models) @pytest.mark.parametrize("ulysses_degree", [1, 2]) @pytest.mark.parametrize("ring_degree", [1, 2]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) # Only test bfloat16 to reduce CI time +@pytest.mark.parametrize("dtype", [torch.bfloat16]) @pytest.mark.parametrize("attn_backend", ["sdpa"]) def test_sequence_parallel( model_name: str, @@ -100,214 +201,60 @@ def test_sequence_parallel( """Compare baseline (ulysses_degree=1) vs SP (ulysses_degree>1) outputs.""" if ulysses_degree <= 1 and ring_degree <= 1: pytest.skip( - "This test compares ulysses_degree * ring_degree = 1 vs ulysses_degree * ring_degree > 1; provide ulysses_degree or ring_degree>1." + "This test compares ulysses_degree * ring_degree = 1 vs ulysses_degree * ring_degree > 1; " + "provide ulysses_degree or ring_degree>1." ) - # Skip if not enough GPUs available for SP run - if device_count() < ulysses_degree * ring_degree: - pytest.skip(f"Test requires {ulysses_degree * ring_degree} GPUs but only {device_count()} available") + sp_size = ulysses_degree * ring_degree + if device_count() < sp_size: + pytest.skip(f"Test requires {sp_size} GPUs but only {device_count()} available") - # Use minimal settings for fast testing height = 256 width = 256 - num_inference_steps = 4 # Minimal steps for fast test seed = 42 - # Step 1: Baseline (no Ulysses sequence parallel) - baseline_parallel_config = DiffusionParallelConfig(ulysses_degree=1, ring_degree=1) - baseline = Omni( - model=model_name, - parallel_config=baseline_parallel_config, - dtype=dtype, - attention_backend=attn_backend, - ) - try: - outputs = baseline.generate( - PROMPT, - height=height, - width=width, - num_inference_steps=num_inference_steps, - guidance_scale=0.0, - generator=torch.Generator(get_device_name()).manual_seed(seed), - num_outputs_per_prompt=1, - ) - baseline_images = outputs[0].request_output[0].images - finally: - baseline.close() - if dist.is_initialized(): - dist.destroy_process_group() - for key in ["MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"]: - os.environ.pop(key, None) - time.sleep(2) # Wait for resources to release + baseline_images = _run_baseline(model_name, dtype, attn_backend, height, width, seed) + assert baseline_images is not None and len(baseline_images) == 1 - assert baseline_images is not None - assert len(baseline_images) == 1 - assert baseline_images[0].width == width - assert baseline_images[0].height == height + sp_images = _run_sp(model_name, dtype, attn_backend, height, width, seed, ulysses_degree, ring_degree) + assert sp_images is not None and len(sp_images) == 1 - # Step 2: SP (Ulysses-SP + Ring-SP) - sp_parallel_config = DiffusionParallelConfig(ulysses_degree=ulysses_degree, ring_degree=ring_degree) - sp = Omni( - model=model_name, - parallel_config=sp_parallel_config, - dtype=dtype, - attention_backend=attn_backend, - ) - try: - outputs = sp.generate( - PROMPT, - height=height, - width=width, - num_inference_steps=num_inference_steps, - guidance_scale=0.0, - generator=torch.Generator(get_device_name()).manual_seed(seed), - num_outputs_per_prompt=1, - ) - sp_images = outputs[0].request_output[0].images - finally: - sp.close() - if dist.is_initialized(): - dist.destroy_process_group() - for key in ["MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"]: - os.environ.pop(key, None) - time.sleep(2) - - assert sp_images is not None - assert len(sp_images) == 1 - assert sp_images[0].width == width - assert sp_images[0].height == height - - # Step 3: Compare outputs mean_abs_diff, max_abs_diff = _diff_metrics(baseline_images[0], sp_images[0]) - # FP16/BF16 may differ slightly due to different computation order under parallelism. - if dtype in (torch.float16, torch.bfloat16): - mean_threshold = 2e-2 - max_threshold = 2e-1 - else: - mean_threshold = 1e-2 - max_threshold = 1e-1 - - print( - "Image diff stats (baseline ulysses_degree*ring_degree=1 vs SP): " - f"mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e}; " - f"thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}; " - f"ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, dtype={dtype}" - ) + mean_threshold = 2e-2 + max_threshold = 2e-1 assert mean_abs_diff <= mean_threshold and max_abs_diff <= max_threshold, ( - f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e} " - f"(thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}); " - f"ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, dtype={dtype}" + f"Image diff exceeded threshold: mean={mean_abs_diff:.6e}, max={max_abs_diff:.6e}" ) @pytest.mark.parametrize("model_name", models) @pytest.mark.parametrize("dtype", [torch.bfloat16]) @pytest.mark.parametrize("attn_backend", ["sdpa"]) -def test_sequence_parallel_ulysses_sp_only( - model_name: str, - dtype: torch.dtype, - attn_backend: str, -): - """Test sequence parallel with ulysses_degree=4, ring_degree=1, and the image size (332x332) where the sequence length is NOT divisible by sp_size.""" +def test_sequence_parallel_ulysses4(model_name: str, dtype: torch.dtype, attn_backend: str): + """Test SP with ulysses_degree=4.""" ulysses_degree = 4 ring_degree = 1 - # Skip if not enough GPUs available for SP run if device_count() < ulysses_degree * ring_degree: pytest.skip(f"Test requires {ulysses_degree * ring_degree} GPUs but only {device_count()} available") - # (272/8) * (272/8) = 17 * 17 = 289 Not divisible by sp_size=4 height = 272 width = 272 - num_inference_steps = 4 # Minimal steps for fast test seed = 42 - # Step 1: Baseline (no Ulysses sequence parallel) - baseline_parallel_config = DiffusionParallelConfig(ulysses_degree=1, ring_degree=1) - baseline = Omni( - model=model_name, - parallel_config=baseline_parallel_config, - dtype=dtype, - attention_backend=attn_backend, - ) - try: - outputs = baseline.generate( - PROMPT, - height=height, - width=width, - num_inference_steps=num_inference_steps, - guidance_scale=0.0, - generator=torch.Generator(get_device_name()).manual_seed(seed), - num_outputs_per_prompt=1, - ) - baseline_images = outputs[0].request_output[0].images - finally: - baseline.close() - if dist.is_initialized(): - dist.destroy_process_group() - for key in ["MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"]: - os.environ.pop(key, None) - time.sleep(2) # Wait for resources to release + baseline_images = _run_baseline(model_name, dtype, attn_backend, height, width, seed) + assert baseline_images is not None and len(baseline_images) == 1 - assert baseline_images is not None - assert len(baseline_images) == 1 - assert baseline_images[0].width == width - assert baseline_images[0].height == height + sp_images = _run_sp(model_name, dtype, attn_backend, height, width, seed, ulysses_degree, ring_degree) + assert sp_images is not None and len(sp_images) == 1 - # Step 2: SP (Ulysses-SP + Ring-SP) - sp_parallel_config = DiffusionParallelConfig(ulysses_degree=ulysses_degree, ring_degree=ring_degree) - sp = Omni( - model=model_name, - parallel_config=sp_parallel_config, - dtype=dtype, - attention_backend=attn_backend, - ) - try: - outputs = sp.generate( - PROMPT, - height=height, - width=width, - num_inference_steps=num_inference_steps, - guidance_scale=0.0, - generator=torch.Generator(get_device_name()).manual_seed(seed), - num_outputs_per_prompt=1, - ) - sp_images = outputs[0].request_output[0].images - finally: - sp.close() - if dist.is_initialized(): - dist.destroy_process_group() - for key in ["MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"]: - os.environ.pop(key, None) - time.sleep(2) - - assert sp_images is not None - assert len(sp_images) == 1 - assert sp_images[0].width == width - assert sp_images[0].height == height - - # Step 3: Compare outputs mean_abs_diff, max_abs_diff = _diff_metrics(baseline_images[0], sp_images[0]) - # FP16/BF16 may differ slightly due to different computation order under parallelism. - if dtype in (torch.float16, torch.bfloat16): - mean_threshold = 2e-2 - max_threshold = 2e-1 - else: - mean_threshold = 1e-2 - max_threshold = 1e-1 - - print( - "Image diff stats (baseline ulysses_degree*ring_degree=1 vs SP): " - f"mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e}; " - f"thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}; " - f"ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, dtype={dtype}" - ) + mean_threshold = 2e-2 + max_threshold = 2e-1 assert mean_abs_diff <= mean_threshold and max_abs_diff <= max_threshold, ( - f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e} " - f"(thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}); " - f"ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, dtype={dtype}" + f"Image diff exceeded threshold: mean={mean_abs_diff:.6e}, max={max_abs_diff:.6e}" ) diff --git a/tools/pre_commit/check_pickle_imports.py b/tools/pre_commit/check_pickle_imports.py index db45f2990..a6c0f3dd3 100644 --- a/tools/pre_commit/check_pickle_imports.py +++ b/tools/pre_commit/check_pickle_imports.py @@ -20,7 +20,7 @@ ALLOWED_FILES = { "tests/e2e/offline_inference/utils.py", "tests/utils.py", "vllm_omni/diffusion/distributed/group_coordinator.py", - "tests/diffusion/attention/test_sequence_parallel.py", + "tests/diffusion/attention/test_attention_sp.py", } PICKLE_RE = re.compile( diff --git a/vllm_omni/diffusion/distributed/__init__.py b/vllm_omni/diffusion/distributed/__init__.py index e69de29bb..c791ae1be 100644 --- a/vllm_omni/diffusion/distributed/__init__.py +++ b/vllm_omni/diffusion/distributed/__init__.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Distributed utilities for vLLM-Omni diffusion models.""" + +from vllm_omni.diffusion.distributed.sp_plan import ( + SequenceParallelConfig, + SequenceParallelInput, + SequenceParallelModelPlan, + SequenceParallelOutput, + SequenceParallelPartialInput, + get_sp_plan_from_model, + validate_sp_plan, +) +from vllm_omni.diffusion.distributed.sp_sharding import ( + ShardingValidator, + get_sharding_validator, + sp_gather, + sp_shard, + sp_shard_with_padding, +) + +__all__ = [ + # Config + "SequenceParallelConfig", + # Plan types + "SequenceParallelInput", + "SequenceParallelOutput", + "SequenceParallelPartialInput", + "SequenceParallelModelPlan", + "validate_sp_plan", + "get_sp_plan_from_model", + # Sharding utilities + "sp_shard", + "sp_gather", + "sp_shard_with_padding", + "ShardingValidator", + "get_sharding_validator", +] diff --git a/vllm_omni/diffusion/distributed/sp_plan.py b/vllm_omni/diffusion/distributed/sp_plan.py new file mode 100644 index 000000000..ee7b1d66b --- /dev/null +++ b/vllm_omni/diffusion/distributed/sp_plan.py @@ -0,0 +1,457 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM and The HuggingFace Team +# Type definitions in this module are adapted from HuggingFace diffusers library: +# diffusers/src/diffusers/models/_modeling_parallel.py +"""Sequence Parallelism configuration and plan type definitions. + +This module provides: +1. SequenceParallelConfig: Configuration for SP (ulysses_degree, ring_degree) +2. SequenceParallelInput/Output: Type definitions for _sp_plan declarations +3. Validation utilities for _sp_plan + +A _sp_plan is a dictionary that specifies how to shard/gather tensors at +different points in a model's forward pass. This allows automatic handling +of sequence parallelism without modifying the model's forward() method. + +NOTE: Our "Sequence Parallelism" (SP) corresponds to "Context Parallelism" (CP) +in diffusers. We use "Sequence Parallelism" to align with vLLM-Omni terminology. + +Example: + class MyTransformer(nn.Module): + _sp_plan = { + # Split inputs before model forward + "": { + "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3), + "encoder_hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3), + }, + # Split RoPE embeddings after pos_embed layer + "pos_embed": { + 0: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True), + }, + # Gather output after proj_out layer + "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), + } +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import torch + import torch.nn as nn + + +# ============================================================================= +# Sequence Parallel Configuration +# ============================================================================= + + +@dataclass +class SequenceParallelConfig: + """Configuration for Sequence Parallelism using vLLM-Omni's parallel state. + + This class provides a unified interface for SP configuration that integrates + with vLLM-Omni's existing SequenceParallelGroupCoordinator. Unlike diffusers' + DeviceMesh-based approach (ContextParallelConfig), this uses the existing + parallel state management. + + Note: This corresponds to `ContextParallelConfig` in diffusers library. + + Args: + ulysses_degree: Number of devices for Ulysses (All-to-All) attention. + Sequence is split across devices, with Q/K/V redistributed via + All-to-All communication. Best for moderate sequences with good + interconnect bandwidth. + ring_degree: Number of devices for Ring attention. Sequence is split + across devices, with K/V passed in a ring topology. Best for long + sequences with limited memory/bandwidth. + convert_to_fp32: Whether to convert output and LSE to float32 for + numerical stability in ring attention. + + Note: + ulysses_degree * ring_degree = sequence_parallel_size + vLLM-Omni supports hybrid Ulysses-Ring attention (both > 1). + """ + + ulysses_degree: int = 1 + ring_degree: int = 1 + convert_to_fp32: bool = True + + # Internal state - populated by setup() + _rank: int | None = None + _world_size: int | None = None + _device: torch.device | None = None + + def __post_init__(self) -> None: + if self.ulysses_degree < 1 or self.ring_degree < 1: + raise ValueError("`ulysses_degree` and `ring_degree` must be >= 1.") + + if self.ulysses_degree == 1 and self.ring_degree == 1: + raise ValueError( + "At least one of `ulysses_degree` or `ring_degree` must be > 1 to use sequence parallelism." + ) + + @property + def sequence_parallel_size(self) -> int: + """Total sequence parallel world size.""" + return self.ulysses_degree * self.ring_degree + + def get_world_size(self) -> int: + """Get the sequence parallel world size from parallel state. + + Returns: + The world size for sequence parallelism. + + Raises: + RuntimeError: If parallel state is not initialized. + """ + from vllm_omni.diffusion.distributed.parallel_state import get_sequence_parallel_world_size + + return get_sequence_parallel_world_size() + + def get_rank(self) -> int: + """Get the current rank in the sequence parallel group. + + Returns: + The rank within the sequence parallel group. + + Raises: + RuntimeError: If parallel state is not initialized. + """ + from vllm_omni.diffusion.distributed.parallel_state import get_sequence_parallel_rank + + return get_sequence_parallel_rank() + + def get_ulysses_world_size(self) -> int: + """Get the Ulysses parallel world size. + + Returns: + The world size for Ulysses (All-to-All) parallelism. + """ + from vllm_omni.diffusion.distributed.parallel_state import get_ulysses_parallel_world_size + + return get_ulysses_parallel_world_size() + + def get_ulysses_rank(self) -> int: + """Get the current rank in the Ulysses parallel group. + + Returns: + The rank within the Ulysses parallel group. + """ + from vllm_omni.diffusion.distributed.parallel_state import get_ulysses_parallel_rank + + return get_ulysses_parallel_rank() + + def get_ring_world_size(self) -> int: + """Get the Ring parallel world size. + + Returns: + The world size for Ring attention parallelism. + """ + from vllm_omni.diffusion.distributed.parallel_state import get_ring_parallel_world_size + + return get_ring_parallel_world_size() + + def get_ring_rank(self) -> int: + """Get the current rank in the Ring parallel group. + + Returns: + The rank within the Ring parallel group. + """ + from vllm_omni.diffusion.distributed.parallel_state import get_ring_parallel_rank + + return get_ring_parallel_rank() + + def setup(self, rank: int, world_size: int, device: torch.device) -> None: + """Initialize the config with runtime parallel state. + + This is called automatically when sequence parallelism is enabled. + + Args: + rank: The global rank of this process. + world_size: Total world size. + device: The device for this rank. + """ + self._rank = rank + self._world_size = world_size + self._device = device + + expected_sp_size = self.ulysses_degree * self.ring_degree + actual_sp_size = self.get_world_size() + + if expected_sp_size != actual_sp_size: + raise ValueError( + f"Configuration mismatch: ulysses_degree ({self.ulysses_degree}) * " + f"ring_degree ({self.ring_degree}) = {expected_sp_size}, but " + f"actual sequence parallel world size is {actual_sp_size}." + ) + + def is_initialized(self) -> bool: + """Check if the config has been initialized with runtime state. + + Returns: + True if setup() has been called, False otherwise. + """ + return self._rank is not None + + +# ============================================================================= +# Sequence Parallel Plan Type Definitions +# ============================================================================= + + +@dataclass(frozen=True) +class SequenceParallelInput: + """Configuration for splitting an input tensor across sequence parallel ranks. + + This specifies how to shard a tensor in the pre-forward or post-forward hook + of a layer. The tensor will be split along the specified dimension. + + Note: This corresponds to `ContextParallelInput` in diffusers library. + + Args: + split_dim: The dimension along which to split the tensor. + expected_dims: Expected number of dimensions. If provided, validates that + the tensor has this many dimensions before splitting. If the tensor + has a different number of dimensions, splitting is skipped with a warning. + split_output: If True, split the output of the layer instead of the input. + This is useful for layers whose outputs should be split after preprocessing + (e.g., RoPE embeddings). + auto_pad: If True, automatically pad the tensor if its size along split_dim + is not divisible by world_size. Creates an attention mask to indicate + valid vs padding positions. The mask is stored in ForwardContext. + Note: Ring attention does not support attention mask, so auto_pad + should only be used with Ulysses SP. + + Example: + # Split hidden_states along sequence dimension (dim 1) + SequenceParallelInput(split_dim=1, expected_dims=3) + + # Split RoPE output along sequence dimension (dim 0) + SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True) + + # Split with auto-padding for variable-length sequences + SequenceParallelInput(split_dim=1, expected_dims=3, auto_pad=True) + """ + + split_dim: int + expected_dims: int | None = None + split_output: bool = False + auto_pad: bool = False + + 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})" + ) + + +@dataclass(frozen=True) +class SequenceParallelOutput: + """Configuration for gathering an output tensor across sequence parallel ranks. + + This specifies how to gather a tensor in the post-forward hook of a layer. + The tensor will be gathered along the specified dimension from all ranks. + + Note: This corresponds to `ContextParallelOutput` in diffusers library. + + Args: + gather_dim: The dimension along which to gather the tensor. + expected_dims: Expected number of dimensions. If provided, validates that + the tensor has this many dimensions before gathering. + + Example: + # Gather output along sequence dimension (dim 1) + SequenceParallelOutput(gather_dim=1, expected_dims=3) + """ + + gather_dim: int + expected_dims: int | None = None + + def __repr__(self) -> str: + return f"SequenceParallelOutput(gather_dim={self.gather_dim}, expected_dims={self.expected_dims})" + + +@dataclass(frozen=True) +class SequenceParallelPartialInput: + """Configuration for partially splitting a tensor (e.g., split image part, keep text part). + + This is designed for models like LongCat/Qwen where RoPE embeddings need special handling: + - Text portion: kept full across all ranks (for joint attention) + - Image portion: split across ranks + + The tensor is assumed to be concatenated as [text_part, image_part] along split_dim. + + Note: This is an extension beyond diffusers' standard ContextParallelInput, + designed for vLLM-Omni's dual-stream attention models. + + Args: + split_dim: The dimension along which to split the image portion. + text_len_source: How to determine text length: + - str: Name of a forward parameter that contains text length + - int: Fixed text length value + expected_dims: Expected number of dimensions for validation. + split_output: If True, split the output instead of input. + + Example: + # Split RoPE: text portion (from txt_ids.shape[0]) kept full, image portion split + SequenceParallelPartialInput( + split_dim=0, + text_len_source="txt_ids", # Get text length from txt_ids.shape[0] + expected_dims=2, + split_output=True, + ) + + # Or with fixed text length + SequenceParallelPartialInput( + split_dim=0, + text_len_source=512, # Fixed text length + expected_dims=2, + split_output=True, + ) + """ + + split_dim: int + text_len_source: str | int + expected_dims: int | None = None + split_output: bool = False + + def __repr__(self) -> str: + return ( + f"SequenceParallelPartialInput(split_dim={self.split_dim}, " + f"text_len_source={self.text_len_source!r}, expected_dims={self.expected_dims}, " + f"split_output={self.split_output})" + ) + + +# ============================================================================= +# Type Aliases for _sp_plan Structure +# ============================================================================= + +# Any input config type +AnySequenceParallelInput = SequenceParallelInput | SequenceParallelPartialInput + +# Input specification: maps parameter names (str) or output indices (int) to split config +SequenceParallelInputType = dict[ + str | int, + AnySequenceParallelInput | list[AnySequenceParallelInput] | tuple[AnySequenceParallelInput, ...], +] + +# Output specification: single or multiple gather configs +SequenceParallelOutputType = SequenceParallelOutput | list[SequenceParallelOutput] | tuple[SequenceParallelOutput, ...] + +# Full model plan: maps module names to input/output specifications +# - Key "" refers to the model itself (root level) +# - Key "module_name" refers to a submodule +# - Key "module_name.*" refers to all children of a ModuleList +# +# Example of a complete _sp_plan: +# +# _sp_plan = { +# # Root level: split model inputs before any submodule +# "": { +# "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3), +# }, +# # Submodule: split outputs of pos_embed (RoPE) layer +# "pos_embed": { +# 0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True), # cos +# 1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True), # sin +# }, +# # Submodule: gather outputs of proj_out layer +# "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), +# } +# +SequenceParallelModelPlan = dict[str, SequenceParallelInputType | SequenceParallelOutputType] + + +# ============================================================================= +# Validation Utilities +# ============================================================================= + + +def _is_valid_input_config(value: object) -> bool: + """Check if a value is a valid input configuration type.""" + return isinstance(value, (SequenceParallelInput, SequenceParallelPartialInput)) + + +def _is_valid_input_config_list(value: object) -> bool: + """Check if a value is a list/tuple of valid input configurations.""" + if not isinstance(value, (list, tuple)): + return False + return all(_is_valid_input_config(x) for x in value) + + +def validate_sp_plan(plan: SequenceParallelModelPlan) -> None: + """Validate a _sp_plan dictionary for correctness. + + Args: + plan: The _sp_plan dictionary to validate. + + Raises: + ValueError: If the plan is invalid. + """ + if not isinstance(plan, dict): + raise ValueError(f"_sp_plan must be a dict, got {type(plan).__name__}") + + for module_id, module_plan in plan.items(): + if not isinstance(module_id, str): + raise ValueError(f"_sp_plan keys must be strings, got {type(module_id).__name__}") + + # Check if it's an output specification (SequenceParallelOutput or list/tuple thereof) + if isinstance(module_plan, SequenceParallelOutput): + continue + if isinstance(module_plan, (list, tuple)): + if all(isinstance(x, SequenceParallelOutput) for x in module_plan): + continue + if _is_valid_input_config_list(module_plan): + # List of inputs for a specific parameter (when output is tuple) + continue + + # Otherwise, should be an input specification dict + if isinstance(module_plan, dict): + for key, value in module_plan.items(): + if not isinstance(key, (str, int)): + raise ValueError( + f"Input spec keys must be str or int, got {type(key).__name__} for module '{module_id}'" + ) + if isinstance(key, int) and not _is_valid_input_config(value): + raise ValueError( + f"Integer keys (output indices) must map to SequenceParallelInput/PartialInput, " + f"got {type(value).__name__} for module '{module_id}'[{key}]" + ) + if _is_valid_input_config(value): + if isinstance(key, int) and not value.split_output: + raise ValueError( + f"Integer keys (output indices) require split_output=True, " + f"got split_output=False for module '{module_id}'[{key}]" + ) + elif _is_valid_input_config_list(value): + pass # Valid list of input configs + else: + raise ValueError( + f"Input spec values must be SequenceParallelInput/PartialInput or list thereof, " + f"got {type(value).__name__} for module '{module_id}'['{key}']" + ) + else: + raise ValueError( + f"_sp_plan values must be dict (input spec) or SequenceParallelOutput, " + f"got {type(module_plan).__name__} for module '{module_id}'" + ) + + +def get_sp_plan_from_model(model: nn.Module) -> SequenceParallelModelPlan | None: + """Get the _sp_plan from a model if it exists. + + Args: + model: The model to get the plan from. + + Returns: + The _sp_plan dictionary, or None if not defined. + """ + plan = getattr(model, "_sp_plan", None) + if plan is not None: + validate_sp_plan(plan) + return plan diff --git a/vllm_omni/diffusion/distributed/sp_sharding.py b/vllm_omni/diffusion/distributed/sp_sharding.py new file mode 100644 index 000000000..12520183c --- /dev/null +++ b/vllm_omni/diffusion/distributed/sp_sharding.py @@ -0,0 +1,268 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project and The HuggingFace Team +"""Sequence Parallelism sharding utilities. + +This module provides low-level sharding and gathering functions for +Sequence Parallelism. These can be used directly in model forward methods +for semi-intrusive SP support, or internally by the SP hooks. +""" + +from __future__ import annotations + +from contextlib import contextmanager +from dataclasses import dataclass, field + +import torch +from vllm.logger import init_logger + +from vllm_omni.diffusion.distributed.parallel_state import ( + get_sequence_parallel_rank, + get_sequence_parallel_world_size, + get_sp_group, +) + +logger = init_logger(__name__) + + +def sp_shard( + tensor: torch.Tensor, + dim: int, + validate: bool = True, +) -> torch.Tensor: + """Shard a tensor along the specified dimension for sequence parallelism. + + The tensor is split into world_size chunks along dim, and this rank + receives its corresponding chunk. + + Args: + tensor: The tensor to shard. + dim: The dimension along which to split. + validate: If True, validate that the tensor size is divisible by world_size. + + Returns: + The shard for this rank. + + Raises: + ValueError: If validate=True and tensor size is not divisible by world_size. + + Example: + # In model forward: + hidden_states = sp_shard(hidden_states, dim=1) + """ + world_size = get_sequence_parallel_world_size() + + if world_size == 1: + return tensor + + rank = get_sequence_parallel_rank() + size = tensor.size(dim) + + if validate and size % world_size != 0: + raise ValueError( + f"Tensor size along dim {dim} ({size}) must be divisible by " + f"world_size ({world_size}) for sequence parallel sharding." + ) + + return tensor.chunk(world_size, dim=dim)[rank] + + +def sp_gather( + tensor: torch.Tensor, + dim: int, + validate: bool = True, +) -> torch.Tensor: + """Gather a tensor along the specified dimension from all sequence parallel ranks. + + The sharded tensors from all ranks are concatenated along dim. + + Args: + tensor: The local shard to gather. + dim: The dimension along which to gather. + validate: If True, validate tensor consistency (currently unused). + + Returns: + The full tensor gathered from all ranks. + + Example: + # At end of model forward: + output = sp_gather(output, dim=1) + """ + world_size = get_sequence_parallel_world_size() + + if world_size == 1: + return tensor + + sp_group = get_sp_group() + return sp_group.all_gather(tensor, dim=dim) + + +def sp_shard_with_padding( + tensor: torch.Tensor, + dim: int, + pad_value: float = 0.0, +) -> tuple[torch.Tensor, int]: + """Shard a tensor with automatic padding if not divisible by world_size. + + This is useful for variable-length sequences where padding may be needed. + + Args: + tensor: The tensor to shard. + dim: The dimension along which to split. + pad_value: Value to use for padding. + + Returns: + Tuple of (sharded_tensor, padding_size). The padding_size indicates + how much padding was added to the original tensor before sharding. + + Example: + sharded, pad_size = sp_shard_with_padding(hidden_states, dim=1) + # ... process ... + output = sp_gather(output, dim=1) + if pad_size > 0: + output = output[..., :-pad_size] # Remove padding + """ + world_size = get_sequence_parallel_world_size() + + if world_size == 1: + return tensor, 0 + + size = tensor.size(dim) + remainder = size % world_size + + if remainder == 0: + return sp_shard(tensor, dim, validate=False), 0 + + # Pad to make divisible + pad_size = world_size - remainder + pad_shape = list(tensor.shape) + pad_shape[dim] = pad_size + padding = torch.full(pad_shape, pad_value, dtype=tensor.dtype, device=tensor.device) + tensor = torch.cat([tensor, padding], dim=dim) + + return sp_shard(tensor, dim, validate=False), pad_size + + +# NOTE: This class is a vLLM-Omni extension for +# debugging intrusive SP implementations. +# Purpose: +# - Help developers detect bugs when implementing intrusive SP +# - Verify that every sharded tensor is properly gathered +# - Warn about common mistakes (double shard, gather without shard) +# +# When to use: +# - During development/debugging of intrusive SP code +# - In tests to verify shard/gather correctness +@dataclass +class ShardingValidator: + """Validator for tracking and verifying sharding operations. + + This class helps ensure that sharding and gathering operations are + correctly paired in model forward passes. It tracks which tensors + have been sharded and verifies that they are properly gathered. + + Usage: + validator = ShardingValidator() + with validator.track(): + hidden_states = validator.shard(hidden_states, "hidden_states", dim=1) + # ... model computation ... + output = validator.gather(output, "hidden_states", dim=1) + validator.validate() # Raises if any shard was not gathered + + Attributes: + _sharded: Set of tensor names that have been sharded. + _gathered: Set of tensor names that have been gathered. + _enabled: Whether tracking is currently enabled. + """ + + _sharded: set[str] = field(default_factory=set) + _gathered: set[str] = field(default_factory=set) + _enabled: bool = False + + def reset(self) -> None: + """Reset the validator state for a new forward pass.""" + self._sharded.clear() + self._gathered.clear() + + @contextmanager + def track(self): + """Context manager to enable tracking for a forward pass.""" + self._enabled = True + self.reset() + try: + yield + finally: + self._enabled = False + + def shard( + self, + tensor: torch.Tensor, + name: str, + dim: int, + validate_divisible: bool = True, + ) -> torch.Tensor: + """Shard a tensor and track the operation. + + Args: + tensor: The tensor to shard. + name: A name to identify this tensor for validation. + dim: The dimension along which to split. + validate_divisible: If True, validate divisibility. + + Returns: + The sharded tensor. + """ + if self._enabled: + if name in self._sharded: + logger.warning(f"Tensor '{name}' sharded multiple times") + self._sharded.add(name) + + return sp_shard(tensor, dim, validate=validate_divisible) + + def gather( + self, + tensor: torch.Tensor, + name: str, + dim: int, + ) -> torch.Tensor: + """Gather a tensor and track the operation. + + Args: + tensor: The local shard to gather. + name: The name used when sharding (for validation). + dim: The dimension along which to gather. + + Returns: + The gathered tensor. + """ + if self._enabled: + if name not in self._sharded: + logger.warning(f"Tensor '{name}' gathered without being sharded") + self._gathered.add(name) + + return sp_gather(tensor, dim) + + def validate(self) -> None: + """Validate that all sharded tensors were gathered. + + Raises: + ValueError: If any sharded tensor was not gathered. + """ + unmatched = self._sharded - self._gathered + if unmatched: + raise ValueError( + f"The following tensors were sharded but not gathered: {unmatched}. " + f"This may indicate a bug in the model's SP implementation." + ) + + +# Global validator instance for convenience +_global_validator = ShardingValidator() + + +def get_sharding_validator() -> ShardingValidator: + """Get the global sharding validator instance. + + Returns: + The global ShardingValidator. + """ + return _global_validator diff --git a/vllm_omni/diffusion/forward_context.py b/vllm_omni/diffusion/forward_context.py index d898bceaf..02082e7b1 100644 --- a/vllm_omni/diffusion/forward_context.py +++ b/vllm_omni/diffusion/forward_context.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from contextlib import contextmanager from dataclasses import dataclass @@ -21,6 +23,14 @@ class ForwardContext: split_text_embed_in_sp: bool = False # whether to split the text embed in sequence parallel, if True, the text embed will be split in sequence parallel + # Sequence Parallel padding support + # When sequence length is not divisible by SP world size, padding is added + # These values are used by SequenceParallelGatherHook to remove padding, + # and by attention layers to create attention masks dynamically + sp_padding_size: int = 0 + # Original sequence length before padding (for removing padding in gather) + sp_original_seq_len: int | None = None + def __post_init__(self): pass diff --git a/vllm_omni/diffusion/hooks.py b/vllm_omni/diffusion/hooks.py deleted file mode 100644 index b1e1016ce..000000000 --- a/vllm_omni/diffusion/hooks.py +++ /dev/null @@ -1,102 +0,0 @@ -from __future__ import annotations - -from collections.abc import Callable -from dataclasses import dataclass -from typing import Any - -import torch.nn as nn - - -class BaseState: - """Base class for hook state containers.""" - - def reset(self) -> None: # pragma: no cover - default is no-op - pass - - -class StateManager: - """Manage per-context hook state instances.""" - - def __init__(self, state_cls: Callable[[], BaseState]): - self._state_cls = state_cls - self._states: dict[str, BaseState] = {} - self._context: str = "default" - - def set_context(self, name: str) -> None: - self._context = name or "default" - - def get_state(self) -> BaseState: - if self._context not in self._states: - self._states[self._context] = self._state_cls() - return self._states[self._context] - - def reset(self) -> None: - self._states.clear() - - -class ModelHook: - """Base class for model hooks that can override a module's forward.""" - - def initialize_hook(self, module: nn.Module) -> nn.Module: - raise NotImplementedError - - def new_forward(self, module: nn.Module, *args: Any, **kwargs: Any): - raise NotImplementedError - - def reset_state(self, module: nn.Module) -> nn.Module: - raise NotImplementedError - - -@dataclass -class _WrappedForward: - module: nn.Module - - def __call__(self, *args: Any, **kwargs: Any): - registry: HookRegistry | None = getattr(self.module, "_hook_registry", None) - if registry is None or not registry._hooks: - return self.module._original_forward(*args, **kwargs) - return registry.dispatch(*args, **kwargs) - - -class HookRegistry: - """Registry of hooks attached to a module.""" - - def __init__(self, module: nn.Module): - self.module = module - self._hooks: dict[str, ModelHook] = {} - - @classmethod - def get_or_create(cls, module: nn.Module) -> HookRegistry: - registry: HookRegistry | None = getattr(module, "_hook_registry", None) - if registry is None: - registry = cls(module) - setattr(module, "_hook_registry", registry) - - # Wrap module.forward once so hooks can intercept calls. - if not hasattr(module, "_original_forward"): - module._original_forward = module.forward # type: ignore[attr-defined] - module.forward = _WrappedForward(module) # type: ignore[assignment] - - return registry - - def register_hook(self, name: str, hook: ModelHook) -> None: - hook.initialize_hook(self.module) - self._hooks[name] = hook - - def get_hook(self, name: str) -> ModelHook | None: - return self._hooks.get(name) - - def dispatch(self, *args: Any, **kwargs: Any): - # For now we support a single active hook and call it directly. - # This can be extended to a chain if needed. - if not self._hooks: - return self.module._original_forward(*args, **kwargs) # type: ignore[attr-defined] - # Deterministic order: sort by name. - name = sorted(self._hooks.keys())[0] - hook = self._hooks[name] - return hook.new_forward(self.module, *args, **kwargs) - - def reset_hook(self, name: str) -> None: - hook = self._hooks.get(name) - if hook is not None: - hook.reset_state(self.module) diff --git a/vllm_omni/diffusion/hooks/__init__.py b/vllm_omni/diffusion/hooks/__init__.py new file mode 100644 index 000000000..68f918383 --- /dev/null +++ b/vllm_omni/diffusion/hooks/__init__.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Hook mechanism for model forward interception.""" + +from vllm_omni.diffusion.hooks.base import ( + BaseState, + HookRegistry, + ModelHook, + StateManager, +) +from vllm_omni.diffusion.hooks.sequence_parallel import ( + SequenceParallelGatherHook, + SequenceParallelSplitHook, + apply_sequence_parallel, + disable_sequence_parallel_for_model, + enable_sequence_parallel_for_model, + remove_sequence_parallel, +) + +__all__ = [ + # Base hooks + "BaseState", + "StateManager", + "ModelHook", + "HookRegistry", + # Sequence parallel hooks (corresponds to diffusers' context_parallel) + "SequenceParallelSplitHook", + "SequenceParallelGatherHook", + "apply_sequence_parallel", + "remove_sequence_parallel", + "enable_sequence_parallel_for_model", + "disable_sequence_parallel_for_model", +] diff --git a/vllm_omni/diffusion/hooks/base.py b/vllm_omni/diffusion/hooks/base.py new file mode 100644 index 000000000..8b330a19f --- /dev/null +++ b/vllm_omni/diffusion/hooks/base.py @@ -0,0 +1,247 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Base hook classes for model forward interception. + +This module provides the foundational hook mechanism that allows intercepting +and modifying model forward passes without invasive changes to model code. +""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +import torch.nn as nn + + +class BaseState: + """Base class for hook state containers.""" + + def reset(self) -> None: # pragma: no cover - default is no-op + pass + + +class StateManager: + """Manage per-context hook state instances.""" + + def __init__(self, state_cls: Callable[[], BaseState]): + self._state_cls = state_cls + self._states: dict[str, BaseState] = {} + self._context: str = "default" + + def set_context(self, name: str) -> None: + self._context = name or "default" + + def get_state(self) -> BaseState: + if self._context not in self._states: + self._states[self._context] = self._state_cls() + return self._states[self._context] + + def reset(self) -> None: + self._states.clear() + + +class ModelHook: + """Base class for model hooks that can override a module's forward. + + Hooks can intercept the forward pass at two points: + - pre_forward: Called before the original forward, can modify args/kwargs + - post_forward: Called after the original forward, can modify output + + Subclasses can override either or both methods. The default implementations + pass through args/kwargs/output unchanged. + + For more complex behavior, override new_forward to completely replace + the forward logic. + """ + + def initialize_hook(self, module: nn.Module) -> nn.Module: + """Initialize the hook when it's registered to a module. + + Args: + module: The module this hook is being attached to. + + Returns: + The module (possibly modified). + """ + return module + + def pre_forward(self, module: nn.Module, *args: Any, **kwargs: Any) -> tuple[tuple, dict]: + """Called before the module's forward pass. + + Args: + module: The module being called. + *args: Positional arguments to forward. + **kwargs: Keyword arguments to forward. + + Returns: + Tuple of (args, kwargs) to pass to the forward method. + """ + return args, kwargs + + def post_forward(self, module: nn.Module, output: Any) -> Any: + """Called after the module's forward pass. + + Args: + module: The module that was called. + output: The output from the forward method. + + Returns: + The (possibly modified) output. + """ + return output + + def new_forward(self, module: nn.Module, *args: Any, **kwargs: Any) -> Any: + """Override the module's forward pass completely. + + The default implementation calls pre_forward, then the original forward, + then post_forward. Override this method for more complex behavior. + + Args: + module: The module being called. + *args: Positional arguments to forward. + **kwargs: Keyword arguments to forward. + + Returns: + The output of the forward pass. + """ + args, kwargs = self.pre_forward(module, *args, **kwargs) + output = module._original_forward(*args, **kwargs) # type: ignore[attr-defined] + return self.post_forward(module, output) + + def reset_state(self, module: nn.Module) -> nn.Module: + """Reset any state associated with this hook. + + Args: + module: The module this hook is attached to. + + Returns: + The module. + """ + return module + + +@dataclass +class _WrappedForward: + """Wrapper that intercepts forward calls and dispatches to hooks.""" + + module: nn.Module + + def __call__(self, *args: Any, **kwargs: Any): + registry: HookRegistry | None = getattr(self.module, "_hook_registry", None) + if registry is None or not registry._hooks: + return self.module._original_forward(*args, **kwargs) + return registry.dispatch(*args, **kwargs) + + +class HookRegistry: + """Registry of hooks attached to a module. + + Manages multiple hooks that can intercept a module's forward pass. + Hooks are called in sorted order by name for determinism. + """ + + def __init__(self, module: nn.Module): + self.module = module + self._hooks: dict[str, ModelHook] = {} + + @classmethod + def get_or_create(cls, module: nn.Module) -> HookRegistry: + """Get existing registry or create a new one for the module. + + Args: + module: The module to get/create a registry for. + + Returns: + The HookRegistry for this module. + """ + registry: HookRegistry | None = getattr(module, "_hook_registry", None) + if registry is None: + registry = cls(module) + setattr(module, "_hook_registry", registry) + + # Wrap module.forward once so hooks can intercept calls. + if not hasattr(module, "_original_forward"): + module._original_forward = module.forward # type: ignore[attr-defined] + module.forward = _WrappedForward(module) # type: ignore[assignment] + + return registry + + def register_hook(self, name: str, hook: ModelHook) -> None: + """Register a hook with the given name. + + Args: + name: Unique name for this hook. + hook: The hook instance to register. + """ + hook.initialize_hook(self.module) + self._hooks[name] = hook + + def remove_hook(self, name: str) -> None: + """Remove a hook by name. + + Args: + name: The name of the hook to remove. + """ + if name in self._hooks: + del self._hooks[name] + + def get_hook(self, name: str) -> ModelHook | None: + """Get a hook by name. + + Args: + name: The name of the hook. + + Returns: + The hook if found, None otherwise. + """ + return self._hooks.get(name) + + def dispatch(self, *args: Any, **kwargs: Any) -> Any: + """Dispatch a forward call through registered hooks. + + Currently supports a single active hook. Multiple hooks are called + in sorted order by name, with each hook's output passed to the next. + + Args: + *args: Positional arguments to forward. + **kwargs: Keyword arguments to forward. + + Returns: + The output of the forward pass. + """ + if not self._hooks: + return self.module._original_forward(*args, **kwargs) # type: ignore[attr-defined] + + # For single hook case, call directly + if len(self._hooks) == 1: + hook = next(iter(self._hooks.values())) + return hook.new_forward(self.module, *args, **kwargs) + + # For multiple hooks, chain them in sorted order + # Each hook can modify args/kwargs via pre_forward + sorted_hooks = sorted(self._hooks.items(), key=lambda x: x[0]) + + # Apply all pre_forward hooks + for _, hook in sorted_hooks: + args, kwargs = hook.pre_forward(self.module, *args, **kwargs) + + # Call original forward + output = self.module._original_forward(*args, **kwargs) # type: ignore[attr-defined] + + # Apply all post_forward hooks in reverse order + for _, hook in reversed(sorted_hooks): + output = hook.post_forward(self.module, output) + + return output + + def reset_hook(self, name: str) -> None: + """Reset a hook's state by name. + + Args: + name: The name of the hook to reset. + """ + hook = self._hooks.get(name) + if hook is not None: + hook.reset_state(self.module) diff --git a/vllm_omni/diffusion/hooks/sequence_parallel.py b/vllm_omni/diffusion/hooks/sequence_parallel.py new file mode 100644 index 000000000..21f4a6fe1 --- /dev/null +++ b/vllm_omni/diffusion/hooks/sequence_parallel.py @@ -0,0 +1,703 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project and the HuggingFace Team. +# All rights reserved. +# +# This module is adapted from HuggingFace diffusers library: +# diffusers/src/diffusers/hooks/context_parallel.py +# +# NOTE: Our "Sequence Parallelism" (SP) corresponds to "Context Parallelism" (CP) in diffusers. +# We use the term "Sequence Parallelism" to align with vLLM-Omni's existing terminology. +# +# Key adaptations from diffusers: +# - ModuleForwardMetadata: parameter lookup logic (adapted) +# - SequenceParallelSplitHook/GatherHook: hook structure (adapted from ContextParallel*) +# - apply_sequence_parallel: registration logic (adapted from apply_context_parallel) +# +# Key differences from diffusers: +# - Uses vLLM-Omni's SequenceParallelGroupCoordinator instead of DeviceMesh +# - Uses sp_shard/sp_gather from sp_sharding.py instead of funcol operations +# - Supports Ulysses + Ring hybrid parallelism +# +"""Sequence Parallelism hooks for non-intrusive SP support. + +This module implements the hook-based mechanism for applying sequence parallelism +to models without modifying their forward() methods. + +Usage: + 1. Define _sp_plan on your model class (corresponds to diffusers' _cp_plan) + 2. Call apply_sequence_parallel(model, config, plan) to enable SP + 3. Call remove_sequence_parallel(model, plan) to disable SP + +The hooks automatically shard inputs before forward and gather outputs after, +based on the plan specification. +""" + +from __future__ import annotations + +import inspect +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn as nn +from vllm.logger import init_logger + +from vllm_omni.diffusion.distributed.sp_plan import ( + AnySequenceParallelInput, + SequenceParallelConfig, + SequenceParallelInput, + SequenceParallelModelPlan, + SequenceParallelOutput, + SequenceParallelPartialInput, +) +from vllm_omni.diffusion.distributed.sp_sharding import sp_gather, sp_shard +from vllm_omni.diffusion.hooks.base import HookRegistry, ModelHook + +logger = init_logger(__name__) + +# Hook name templates for identifying SP hooks +_SP_INPUT_HOOK_TEMPLATE = "sp_input---{}" +_SP_OUTPUT_HOOK_TEMPLATE = "sp_output---{}" + + +@dataclass +class ModuleForwardMetadata: + """Metadata for mapping forward() parameter names to args/kwargs positions. + + This caches the inspection of a module's forward signature to efficiently + locate parameters by name in subsequent calls. + """ + + cached_parameter_indices: dict[str, int] | None = None + _cls: type | None = None + + def _get_parameter_from_args_kwargs( + self, + identifier: str, + args: tuple = (), + kwargs: dict | None = None, + ) -> tuple[Any, bool, int | None]: + """Get a parameter value from args or kwargs by name. + + Args: + identifier: The parameter name to look up. + args: Positional arguments passed to forward. + kwargs: Keyword arguments passed to forward. + + Returns: + Tuple of (value, is_kwarg, index). + - value: The parameter value (or None if not found) + - is_kwarg: True if found in kwargs + - index: Position in args if found there + + Raises: + ValueError: If parameter not found in signature. + """ + kwargs = kwargs or {} + + # First check kwargs + if identifier in kwargs: + return kwargs[identifier], True, None + + # Check cached indices + if self.cached_parameter_indices is not None: + index = self.cached_parameter_indices.get(identifier, None) + if index is None: + raise ValueError(f"Parameter '{identifier}' not found in cached indices.") + if index < len(args): + return args[index], False, index + return None, False, index + + # Build cache from forward signature + if self._cls is None: + raise ValueError("Model class is not set for metadata.") + + parameters = list(inspect.signature(self._cls.forward).parameters.keys()) + parameters = parameters[1:] # Skip `self` + self.cached_parameter_indices = {param: i for i, param in enumerate(parameters)} + + if identifier not in self.cached_parameter_indices: + raise ValueError(f"Parameter '{identifier}' not found in function signature.") + + index = self.cached_parameter_indices[identifier] + + if index >= len(args): + return None, False, index + + return args[index], False, index + + +def _unwrap_module(module: nn.Module) -> nn.Module: + """Unwrap a module from any wrappers to get the original class. + + Args: + module: Potentially wrapped module. + + Returns: + The unwrapped module. + """ + # Handle common wrappers + while hasattr(module, "_modules") and len(module._modules) == 1: + inner = next(iter(module._modules.values())) + if inner is not None: + module = inner + else: + break + return module + + +class SequenceParallelSplitHook(ModelHook): + """Hook for splitting inputs before a module's forward pass. + + This hook is registered to modules that need their inputs sharded + across sequence parallel ranks. It intercepts the forward call, + shards specified inputs according to the plan, and passes the + sharded inputs to the original forward. + + For split_output=True inputs, it shards the output instead. + + Supports both SequenceParallelInput (full split) and SequenceParallelPartialInput + (partial split for text/image separation). + + Note: This corresponds to `ContextParallelSplitHook` in diffusers. + """ + + def __init__( + self, + metadata: dict[str | int, AnySequenceParallelInput | list[AnySequenceParallelInput]], + config: SequenceParallelConfig, + ) -> None: + super().__init__() + self.metadata = metadata + self.config = config + self.module_forward_metadata: ModuleForwardMetadata | None = None + # Cache for text lengths resolved from kwargs + self._text_len_cache: dict[str, int] = {} + + def initialize_hook(self, module: nn.Module) -> nn.Module: + cls = _unwrap_module(module).__class__ + self.module_forward_metadata = ModuleForwardMetadata(_cls=cls) + return module + + def pre_forward(self, module: nn.Module, *args: Any, **kwargs: Any) -> tuple[tuple, dict]: + """Shard inputs before forward.""" + args_list = list(args) + # Clear text length cache for this forward pass + self._text_len_cache.clear() + + for name, spm in self.metadata.items(): + # Skip if this is a split_output entry (handled in post_forward) + if isinstance(spm, (SequenceParallelInput, SequenceParallelPartialInput)) and spm.split_output: + continue + + # Get the parameter value + input_val, is_kwarg, index = self.module_forward_metadata._get_parameter_from_args_kwargs( + name, args_list, kwargs + ) + + if input_val is None: + continue + + # Shard the input + if isinstance(input_val, torch.Tensor): + input_val = self._prepare_sp_input(input_val, spm, args_list, kwargs) + elif isinstance(input_val, (list, tuple)): + # Handle list/tuple of tensors with per-element config + if not isinstance(spm, (list, tuple)): + raise ValueError( + f"Expected list/tuple of SequenceParallelInput for parameter '{name}' " + f"which is a list/tuple, but got {type(spm).__name__}" + ) + if len(input_val) != len(spm): + raise ValueError(f"Expected {len(spm)} elements for parameter '{name}', got {len(input_val)}") + sharded_input_val = [] + for i, x in enumerate(input_val): + if torch.is_tensor(x) and not spm[i].split_output: + x = self._prepare_sp_input(x, spm[i], args_list, kwargs) + sharded_input_val.append(x) + input_val = type(input_val)(sharded_input_val) + else: + raise ValueError(f"Unsupported input type for sharding: {type(input_val).__name__}") + + # Update args or kwargs + if is_kwarg: + kwargs[name] = input_val + elif index is not None and index < len(args_list): + args_list[index] = input_val + 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) + + return tuple(args_list), kwargs + + def post_forward(self, module: nn.Module, output: Any) -> Any: + """Shard outputs for split_output=True entries.""" + is_tensor = isinstance(output, torch.Tensor) + is_tensor_list = isinstance(output, (list, tuple)) and all(isinstance(x, torch.Tensor) for x in output) + + if not is_tensor and not is_tensor_list: + # No tensor outputs to shard + return output + + output_list = [output] if is_tensor else list(output) + + for index, spm in self.metadata.items(): + if not isinstance(index, int): + continue + if not isinstance(spm, (SequenceParallelInput, SequenceParallelPartialInput)) or not spm.split_output: + continue + if index >= len(output_list): + raise ValueError(f"Index {index} out of bounds for output of length {len(output_list)}.") + + output_list[index] = self._prepare_sp_input(output_list[index], spm, self._last_args, self._last_kwargs) + + return output_list[0] if is_tensor else type(output)(output_list) + + def _resolve_text_len( + self, + sp_input: SequenceParallelPartialInput, + args: tuple, + kwargs: dict, + ) -> int: + """Resolve text length from the source specification.""" + source = sp_input.text_len_source + + if isinstance(source, int): + return source + + # String source - look up from kwargs or args + if source in self._text_len_cache: + return self._text_len_cache[source] + + # Try to get from kwargs/args + try: + val, _, _ = self.module_forward_metadata._get_parameter_from_args_kwargs(source, args, kwargs) + if val is None: + raise ValueError(f"Parameter '{source}' is None, cannot determine text length.") + if isinstance(val, torch.Tensor): + # TODO: Currently assumes batch_size=1, where shape[0] is sequence length. + # For batch inference support, this should be updated to handle + # shape (batch_size, seq_len, ...) where text_len varies per sample. + text_len = val.shape[0] + elif isinstance(val, int): + text_len = val + else: + raise ValueError(f"Cannot determine text length from '{source}' of type {type(val).__name__}") + self._text_len_cache[source] = text_len + return text_len + except ValueError as e: + raise ValueError(f"Failed to resolve text_len_source '{source}': {e}") from e + + def _prepare_sp_input( + self, + x: torch.Tensor, + sp_input: AnySequenceParallelInput, + args: tuple = (), + kwargs: dict | None = None, + ) -> torch.Tensor: + """Shard a tensor according to the input specification.""" + kwargs = kwargs or {} + + if sp_input.expected_dims is not None and x.dim() != sp_input.expected_dims: + logger.warning_once(f"Expected tensor with {sp_input.expected_dims} dims, got {x.dim()}. Skipping split.") + return x + + if isinstance(sp_input, SequenceParallelInput): + # Full split with optional auto-padding + if sp_input.auto_pad: + return self._shard_with_auto_pad(x, sp_input.split_dim) + return sp_shard(x, sp_input.split_dim, validate=False) + elif isinstance(sp_input, SequenceParallelPartialInput): + # Partial split: keep text portion, split image portion + text_len = self._resolve_text_len(sp_input, args, kwargs) + dim = sp_input.split_dim + + # Split tensor into text and image portions + text_part = x.narrow(dim, 0, text_len) + image_part = x.narrow(dim, text_len, x.size(dim) - text_len) + + # Only shard the image portion + image_part_sharded = sp_shard(image_part, dim, validate=False) + + # Concatenate back: [text_full, image_sharded] + return torch.cat([text_part, image_part_sharded], dim=dim) + else: + raise ValueError(f"Unsupported input config type: {type(sp_input).__name__}") + + def _shard_with_auto_pad(self, x: torch.Tensor, dim: int) -> torch.Tensor: + """Shard tensor with automatic padding and attention mask creation. + + When sequence length is not divisible by SP world size, this method: + 1. Pads the tensor to make it divisible + 2. Creates an attention mask indicating valid vs padding positions + 3. Stores the mask and padding info in ForwardContext + """ + from vllm_omni.diffusion.attention.selector import ( + _BACKENDS_SUPPORT_ATTENTION_MASK, + get_attn_backend, + ) + from vllm_omni.diffusion.distributed.parallel_state import ( + get_ring_parallel_world_size, + get_sequence_parallel_rank, + get_sequence_parallel_world_size, + ) + from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available + + world_size = get_sequence_parallel_world_size() + if world_size == 1: + return x + + seq_len = x.size(dim) + remainder = seq_len % world_size + + if remainder == 0: + # No padding needed + return sp_shard(x, dim, validate=False) + + # Check backend compatibility + attn_backend = get_attn_backend(-1) + if attn_backend.get_name() not in _BACKENDS_SUPPORT_ATTENTION_MASK: + raise ValueError( + f"Sequence length ({seq_len}) is not divisible by SP world size ({world_size}). " + f"Cannot use {attn_backend.get_name()} which does not support attention_mask. " + f"Please switch to SDPA or Ascend attention backend." + ) + + # Ring attention does not support attention_mask + if get_ring_parallel_world_size() > 1: + raise ValueError( + f"Sequence length ({seq_len}) is not divisible by SP world size ({world_size}). " + f"Cannot use Ring attention which does not support attention_mask. " + f"Please switch to Ulysses SP only." + ) + + # Calculate padding + pad_size = world_size - remainder + padded_seq_len = seq_len + pad_size + + # Pad the tensor + pad_shape = list(x.shape) + pad_shape[dim] = pad_size + padding = torch.zeros(pad_shape, dtype=x.dtype, device=x.device) + x_padded = torch.cat([x, padding], dim=dim) + + # Store padding info in forward context (only once, for primary tensor) + # Attention layers will create masks dynamically using this info + if is_forward_context_available(): + ctx = get_forward_context() + # Only set if not already set (first auto_pad tensor wins) + if ctx.sp_original_seq_len is None: + ctx.sp_padding_size = pad_size + ctx.sp_original_seq_len = seq_len + logger.debug( + f"Auto-padded sequence from {seq_len} to {padded_seq_len} " + f"(pad_size={pad_size}, world_size={world_size}, dim={dim})" + ) + + # Shard the padded tensor + rank = get_sequence_parallel_rank() + return x_padded.chunk(world_size, dim=dim)[rank] + + +class SequenceParallelGatherHook(ModelHook): + """Hook for gathering outputs after a module's forward pass. + + This hook is registered to modules that need their outputs gathered + from all sequence parallel ranks. It intercepts the output and gathers + it according to the plan specification. + + Note: This corresponds to `ContextParallelGatherHook` in diffusers. + """ + + def __init__( + self, + metadata: SequenceParallelOutput | list[SequenceParallelOutput], + config: SequenceParallelConfig, + ) -> None: + super().__init__() + if isinstance(metadata, SequenceParallelOutput): + metadata = [metadata] + self.metadata = metadata + self.config = config + + def initialize_hook(self, module: nn.Module) -> nn.Module: + return module + + def post_forward(self, module: nn.Module, output: Any) -> Any: + """Gather outputs after forward and remove padding if applied.""" + from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available + + is_tensor = isinstance(output, torch.Tensor) + + if is_tensor: + output = [output] + elif not (isinstance(output, (list, tuple)) and all(isinstance(x, torch.Tensor) for x in output)): + raise ValueError(f"Expected tensor or list/tuple of tensors, got {type(output).__name__}") + + output = list(output) + + if len(output) != len(self.metadata): + raise ValueError(f"Expected {len(self.metadata)} outputs, got {len(output)}.") + + # Check if padding was applied during split + original_seq_len = None + if is_forward_context_available(): + ctx = get_forward_context() + original_seq_len = ctx.sp_original_seq_len + + for i, spm in enumerate(self.metadata): + if spm is None: + continue + + x = output[i] + if spm.expected_dims is not None and x.dim() != spm.expected_dims: + logger.warning_once( + f"Expected output tensor with {spm.expected_dims} dims, got {x.dim()}. Skipping gather." + ) + continue + + # Gather from all ranks + gathered = sp_gather(x, spm.gather_dim, validate=False) + + # Remove padding if it was applied + if original_seq_len is not None and gathered.size(spm.gather_dim) > original_seq_len: + gathered = gathered.narrow(spm.gather_dim, 0, original_seq_len) + logger.debug(f"Removed padding: gathered shape {gathered.shape} (original_seq_len={original_seq_len})") + + output[i] = gathered + + return output[0] if is_tensor else type(output)(output) + + +def _get_submodule_by_name(model: nn.Module, name: str) -> nn.Module | list[nn.Module]: + """Get a submodule by dotted name, supporting wildcards. + + Args: + model: The root module. + name: Dotted path to submodule. Use "*" to match all children + of a ModuleList. + + Returns: + The submodule or list of submodules if wildcard used. + + Raises: + ValueError: If the path is invalid or module not found. + """ + if name.count("*") > 1: + raise ValueError("Wildcard '*' can only be used once in the name") + return _find_submodule_by_name(model, name) + + +def _find_submodule_by_name(model: nn.Module, name: str) -> nn.Module | list[nn.Module]: + """Recursive helper for _get_submodule_by_name.""" + if name == "": + return model + + first_atom, remaining_name = name.split(".", 1) if "." in name else (name, "") + + if first_atom == "*": + if not isinstance(model, nn.ModuleList): + raise ValueError("Wildcard '*' can only be used with ModuleList") + submodules = [] + for submodule in model: + subsubmodules = _find_submodule_by_name(submodule, remaining_name) + if not isinstance(subsubmodules, list): + subsubmodules = [subsubmodules] + submodules.extend(subsubmodules) + return submodules + else: + if hasattr(model, first_atom): + submodule = getattr(model, first_atom) + return _find_submodule_by_name(submodule, remaining_name) + else: + raise ValueError(f"'{first_atom}' is not a submodule of '{model.__class__.__name__}'") + + +def apply_sequence_parallel( + module: nn.Module, + config: SequenceParallelConfig, + plan: SequenceParallelModelPlan, +) -> None: + """Apply sequence parallel hooks to a model according to the plan. + + This function registers hooks on the specified submodules to automatically + shard inputs and gather outputs for sequence parallelism. + + Note: This corresponds to `apply_context_parallel` in diffusers. + + The complete SP flow is: + 1. Input sharding (SequenceParallelSplitHook): Split sequence across SP ranks + 2. Attention parallelism (handled by vLLM-Omni's Attention layer): + - Ulysses: All-to-All over Q/K/V heads + - Ring: K/V circulation in ring topology + - Hybrid: Both (Ulysses handles head redistribution, Ring handles K/V) + 3. Output gathering (SequenceParallelGatherHook): Gather sequence from SP ranks + + Args: + module: The model to apply SP to. + config: The sequence parallel configuration. + plan: Dictionary mapping module names to input/output specifications. + + Example: + config = SequenceParallelConfig(ulysses_degree=2) + plan = { + "": {"hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3)}, + "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), + } + apply_sequence_parallel(model, config, plan) + + Note: + vLLM-Omni's Attention layer automatically handles the internal + parallelism (Ulysses All-to-All or Ring attention) based on the + forward_context configuration. This function only handles input/output + sharding for the model as a whole. + """ + logger.debug( + f"Applying sequence parallel with config: ulysses={config.ulysses_degree}, " + f"ring={config.ring_degree}, plan keys: {list(plan.keys())}" + ) + + for module_id, sp_model_plan in plan.items(): + submodule = _get_submodule_by_name(module, module_id) + if not isinstance(submodule, list): + submodule = [submodule] + + logger.debug(f"Applying SP hooks to '{module_id}' ({len(submodule)} module(s))") + + for m in submodule: + if isinstance(sp_model_plan, dict): + # Input specification + hook = SequenceParallelSplitHook(sp_model_plan, config) + hook_name = _SP_INPUT_HOOK_TEMPLATE.format(module_id) + elif isinstance(sp_model_plan, (SequenceParallelOutput, list, tuple)): + # Output specification + if isinstance(sp_model_plan, SequenceParallelOutput): + sp_model_plan = [sp_model_plan] + if not all(isinstance(x, SequenceParallelOutput) or x is None for x in sp_model_plan): + raise ValueError(f"Expected SequenceParallelOutput elements, got {sp_model_plan}") + hook = SequenceParallelGatherHook(sp_model_plan, config) + hook_name = _SP_OUTPUT_HOOK_TEMPLATE.format(module_id) + else: + raise ValueError(f"Unsupported plan type: {type(sp_model_plan).__name__}") + + registry = HookRegistry.get_or_create(m) + registry.register_hook(hook_name, hook) + + +def remove_sequence_parallel( + module: nn.Module, + plan: SequenceParallelModelPlan, +) -> None: + """Remove sequence parallel hooks from a model. + + Note: This corresponds to `remove_context_parallel` in diffusers. + + Args: + module: The model to remove SP from. + plan: The same plan used when applying SP. + """ + for module_id, sp_model_plan in plan.items(): + submodule = _get_submodule_by_name(module, module_id) + if not isinstance(submodule, list): + submodule = [submodule] + + for m in submodule: + registry = getattr(m, "_hook_registry", None) + if registry is None: + continue + + if isinstance(sp_model_plan, dict): + hook_name = _SP_INPUT_HOOK_TEMPLATE.format(module_id) + elif isinstance(sp_model_plan, (SequenceParallelOutput, list, tuple)): + hook_name = _SP_OUTPUT_HOOK_TEMPLATE.format(module_id) + else: + continue + + registry.remove_hook(hook_name) + + +def enable_sequence_parallel_for_model( + model: nn.Module, + config: SequenceParallelConfig | None = None, +) -> None: + """Enable sequence parallelism for a model using its _sp_plan. + + This is a convenience function that reads the model's _sp_plan attribute + and applies sequence parallelism automatically. + + Note: This corresponds to `enable_context_parallel_for_model` in diffusers, + but uses vLLM-Omni's _sp_plan instead of diffusers' _cp_plan. + + The function performs two main tasks: + 1. Applies _sp_plan hooks to shard inputs and gather outputs + 2. Ensures Attention layers are configured for the correct parallel mode + (handled automatically by vLLM-Omni's forward_context mechanism) + + Args: + model: The model to enable SP for. Must have a _sp_plan attribute. + config: Optional config. If None, uses default based on current + parallel state. + + Raises: + ValueError: If model has no _sp_plan defined. + + Note: + vLLM-Omni supports Ulysses + Ring hybrid parallelism: + - ulysses_degree > 1: Uses All-to-All communication over Q/K/V heads + - ring_degree > 1: Uses Ring attention with K/V passing + - Both > 1: Hybrid mode (Ulysses handles head redistribution, + Ring handles K/V circulation) + """ + from vllm_omni.diffusion.distributed.parallel_state import ( + get_ring_parallel_world_size, + get_ulysses_parallel_world_size, + ) + from vllm_omni.diffusion.distributed.sp_plan import get_sp_plan_from_model + + plan = get_sp_plan_from_model(model) + if plan is None: + raise ValueError( + f"Model {model.__class__.__name__} has no _sp_plan defined. " + f"Define _sp_plan as a class attribute or pass a plan explicitly." + ) + + if config is None: + # Create config from current parallel state + ulysses_degree = get_ulysses_parallel_world_size() + ring_degree = get_ring_parallel_world_size() + config = SequenceParallelConfig( + ulysses_degree=ulysses_degree, + ring_degree=ring_degree, + ) + if ulysses_degree > 1 and ring_degree > 1: + mode = "hybrid" + elif ulysses_degree > 1: + mode = "ulysses" + else: + mode = "ring" + logger.info( + f"Created SP config from parallel state: " + f"ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, " + f"mode={mode}" + ) + + apply_sequence_parallel(model, config, plan) + logger.info(f"Enabled sequence parallelism for {model.__class__.__name__}") + + +def disable_sequence_parallel_for_model(model: nn.Module) -> None: + """Disable sequence parallelism for a model. + + Note: This corresponds to `disable_context_parallel_for_model` in diffusers. + + Args: + model: The model to disable SP for. + """ + from vllm_omni.diffusion.distributed.sp_plan import get_sp_plan_from_model + + plan = get_sp_plan_from_model(model) + if plan is not None: + remove_sequence_parallel(model, plan) diff --git a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py index cbf0b7e10..14905e2a8 100644 --- a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py +++ b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py @@ -27,14 +27,11 @@ from vllm_omni.diffusion.attention.backends.abstract import ( AttentionMetadata, ) from vllm_omni.diffusion.attention.layer import Attention -from vllm_omni.diffusion.attention.selector import _BACKENDS_SUPPORT_ATTENTION_MASK, get_attn_backend from vllm_omni.diffusion.cache.base import CachedTransformer from vllm_omni.diffusion.data import OmniDiffusionConfig -from vllm_omni.diffusion.distributed.parallel_state import ( - get_ring_parallel_world_size, - get_sequence_parallel_rank, - get_sequence_parallel_world_size, - get_sp_group, +from vllm_omni.diffusion.distributed.sp_plan import ( + SequenceParallelInput, + SequenceParallelOutput, ) from vllm_omni.diffusion.forward_context import get_forward_context from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm @@ -43,6 +40,55 @@ from vllm_omni.diffusion.layers.rope import RotaryEmbedding logger = init_logger(__name__) +class ImageRopePrepare(nn.Module): + """Prepares image hidden_states and RoPE embeddings for sequence parallel. + + This module encapsulates the input linear projection and RoPE computation. + Similar to Z-Image's UnifiedPrepare, this creates a module boundary where + _sp_plan can shard outputs via split_output=True. + + The key insight is that hidden_states and vid_freqs must be sharded together + to maintain dimension alignment for RoPE computation in attention layers. + + Note: Our _sp_plan corresponds to diffusers' _cp_plan (Context Parallelism). + """ + + def __init__(self, img_in: nn.Linear, pos_embed: nn.Module): + super().__init__() + self.img_in = img_in + self.pos_embed = pos_embed + + def forward( + self, + hidden_states: torch.Tensor, + img_shapes: list[tuple[int, int, int]], + txt_seq_lens: list[int], + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Prepare hidden_states and RoPE for SP. + + Args: + hidden_states: [batch, img_seq_len, channels] + img_shapes: List of (frame, height, width) tuples + txt_seq_lens: List of text sequence lengths + + Returns: + hidden_states: Processed hidden states [batch, img_seq_len, dim] + vid_freqs: Image RoPE frequencies [img_seq_len, rope_dim] + txt_freqs: Text RoPE frequencies [txt_seq_len, rope_dim] + + Note: _sp_plan will shard hidden_states and vid_freqs via split_output=True + txt_freqs is kept replicated for dual-stream attention + """ + # Apply input projection + hidden_states = self.img_in(hidden_states) + + # Compute RoPE embeddings + image_rotary_emb = self.pos_embed(img_shapes, txt_seq_lens, device=hidden_states.device) + vid_freqs, txt_freqs = image_rotary_emb + + return hidden_states, vid_freqs, txt_freqs + + class QwenTimestepProjEmbeddings(nn.Module): def __init__(self, embedding_dim, use_additional_t_cond=False): super().__init__() @@ -436,6 +482,25 @@ class QwenImageCrossAttention(nn.Module): hidden_states_mask: torch.Tensor | None = None, encoder_hidden_states_mask: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: + # Check for SP auto_pad: create attention mask dynamically if padding was applied + # In Ulysses mode, attention is computed on the FULL sequence (after All-to-All) + if ( + self.parallel_config is not None + and self.parallel_config.sequence_parallel_size > 1 + and hidden_states_mask is None + ): + ctx = get_forward_context() + if ctx.sp_original_seq_len is not None and ctx.sp_padding_size > 0: + # Create mask for the full (padded) sequence + # valid positions = True, padding positions = False + batch_size = hidden_states.shape[0] + padded_seq_len = ctx.sp_original_seq_len + ctx.sp_padding_size + hidden_states_mask = torch.ones( + batch_size, padded_seq_len, dtype=torch.bool, device=hidden_states.device + ) + hidden_states_mask[:, ctx.sp_original_seq_len :] = False + + # if mask is all true, set it to None if hidden_states_mask is not None and hidden_states_mask.all(): hidden_states_mask = None if encoder_hidden_states_mask is not None and encoder_hidden_states_mask.all(): @@ -718,6 +783,31 @@ class QwenImageTransformer2DModel(CachedTransformer): # used for torch compile optimizations _repeated_blocks = ["QwenImageTransformerBlock"] + # Sequence Parallelism plan (following diffusers' _cp_plan pattern) + # Similar to Z-Image's UnifiedPrepare, we use ImageRopePrepare to create + # a module boundary where _sp_plan can shard hidden_states and vid_freqs together. + # + # Key insight: hidden_states and vid_freqs MUST be sharded together to maintain + # dimension alignment for RoPE computation in attention layers. + # + # auto_pad=True enables automatic padding when sequence length is not divisible + # by SP world size. This creates an attention mask stored in ForwardContext + # that attention layers can use to ignore padding positions. + # + # Note: _sp_plan corresponds to diffusers' _cp_plan (Context Parallelism) + _sp_plan = { + # Shard ImageRopePrepare outputs (hidden_states and vid_freqs must be sharded together) + "image_rope_prepare": { + # hidden_states: auto_pad=True for variable sequence length support + 0: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True, auto_pad=True), + # vid_freqs: auto_pad=True to match hidden_states padding + 1: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True, auto_pad=True), + # txt_freqs (index 2) is NOT sharded - kept replicated for dual-stream attention + }, + # Gather output at proj_out + "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), + } + def __init__( self, od_config: OmniDiffusionConfig, @@ -773,6 +863,10 @@ class QwenImageTransformer2DModel(CachedTransformer): self.gradient_checkpointing = False self.zero_cond_t = zero_cond_t + # ImageRopePrepare module for _sp_plan to shard hidden_states and vid_freqs together + # This ensures RoPE dimensions align with hidden_states after sharding + self.image_rope_prepare = ImageRopePrepare(self.img_in, self.pos_embed) + def forward( self, hidden_states: torch.Tensor, @@ -816,49 +910,17 @@ class QwenImageTransformer2DModel(CachedTransformer): # else: # lora_scale = 1.0 - original_seq_len = None - seq_padding = 0 - hidden_states_mask = None - + # Set split_text_embed_in_sp = False for dual-stream attention + # QwenImage uses *dual-stream* (text + image) and runs a *joint attention*. + # Text embeddings must be replicated across SP ranks for correctness. if self.parallel_config.sequence_parallel_size > 1: - batch_size, seq_len, channels = hidden_states.shape - sp_size = get_sequence_parallel_world_size() - - if seq_len % sp_size != 0: - # ring_attn, sage_attn do not support attention_mask - if get_attn_backend(-1).get_name() not in _BACKENDS_SUPPORT_ATTENTION_MASK: - raise ValueError( - f"When generating image shape that the sequence length is NOT divisible by sp_size={sp_size}," - f"cannot use {get_attn_backend(-1).get_name()} which does not support attention_mask." - f"Please switch to SDPA or Ascend attention backend." - ) - # ring attention does not support attention_mask - if get_ring_parallel_world_size() > 1: - raise ValueError( - f"When generating image shape that the sequence length is NOT divisible by sp_size={sp_size}," - f"cannot use ring attention which does not support attention_mask." - f"Please switch to Ulysses SP only." - ) - - seq_padding = sp_size - (seq_len % sp_size) - original_seq_len = seq_len - - hidden_states_mask = torch.ones( - batch_size, seq_len + seq_padding, dtype=torch.bool, device=hidden_states.device - ) - hidden_states_mask[:, seq_len:] = False - padding_tensor = torch.zeros( - batch_size, seq_padding, channels, dtype=hidden_states.dtype, device=hidden_states.device - ) - hidden_states = torch.cat([hidden_states, padding_tensor], dim=1) - - hidden_states = torch.chunk(hidden_states, sp_size, dim=-2)[get_sequence_parallel_rank()] - # NOTE: - # QwenImage uses *dual-stream* (text + image) and runs a *joint attention*. - # text embeddings to be replicated across SP ranks for correctness. get_forward_context().split_text_embed_in_sp = False - hidden_states = self.img_in(hidden_states) + # Prepare hidden_states and RoPE via ImageRopePrepare module + # _sp_plan will shard hidden_states and vid_freqs together via split_output=True + # txt_freqs is kept replicated for dual-stream attention + hidden_states, vid_freqs, txt_freqs = self.image_rope_prepare(hidden_states, img_shapes, txt_seq_lens) + image_rotary_emb = (vid_freqs, txt_freqs) # Ensure timestep tensor is on the same device and dtype as hidden_states timestep = timestep.to(device=hidden_states.device, dtype=hidden_states.dtype) @@ -885,23 +947,6 @@ class QwenImageTransformer2DModel(CachedTransformer): else self.time_text_embed(timestep, guidance, hidden_states, additional_t_cond) ) - image_rotary_emb = self.pos_embed(img_shapes, txt_seq_lens, device=hidden_states.device) - - def get_rotary_emb_chunk(freqs, padding=0): - # Pad rotary embeddings if needed - if padding > 0: - padding_tensor = torch.zeros(padding, freqs.shape[-1], dtype=freqs.dtype, device=freqs.device) - freqs = torch.cat([freqs, padding_tensor], dim=0) - freqs = torch.chunk(freqs, get_sequence_parallel_world_size(), dim=0)[get_sequence_parallel_rank()] - return freqs - - if self.parallel_config.sequence_parallel_size > 1: - img_freqs, txt_freqs = image_rotary_emb - img_freqs = get_rotary_emb_chunk(img_freqs, seq_padding) - if get_forward_context().split_text_embed_in_sp: - txt_freqs = get_rotary_emb_chunk(txt_freqs) - image_rotary_emb = (img_freqs, txt_freqs) - for index_block, block in enumerate(self.transformer_blocks): encoder_hidden_states, hidden_states = block( hidden_states=hidden_states, @@ -911,7 +956,6 @@ class QwenImageTransformer2DModel(CachedTransformer): image_rotary_emb=image_rotary_emb, joint_attention_kwargs=attention_kwargs, modulate_index=modulate_index, - hidden_states_mask=hidden_states_mask, ) if self.zero_cond_t: @@ -920,14 +964,8 @@ class QwenImageTransformer2DModel(CachedTransformer): hidden_states = self.norm_out(hidden_states, temb) output = self.proj_out(hidden_states) - if self.parallel_config.sequence_parallel_size > 1: - output = get_sp_group().all_gather(output, dim=-2) - - # Remove padding if it was added - if original_seq_len is not None: - output = output[:, :original_seq_len, :] - - torch.cuda.empty_cache() + # Note: SP gather is handled automatically by _sp_plan's SequenceParallelGatherHook + # on proj_out output. No manual all_gather needed here. return Transformer2DModelOutput(sample=output) diff --git a/vllm_omni/diffusion/models/z_image/pipeline_z_image.py b/vllm_omni/diffusion/models/z_image/pipeline_z_image.py index 81539808f..7d6fb901d 100644 --- a/vllm_omni/diffusion/models/z_image/pipeline_z_image.py +++ b/vllm_omni/diffusion/models/z_image/pipeline_z_image.py @@ -176,6 +176,9 @@ class ZImagePipeline(nn.Module): self.transformer = ZImageTransformer2DModel() self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) + # Note: Context parallelism is applied centrally in registry.initialize_model() + # following diffusers' pattern of enable_parallelism() at model loading time + self.vae_scale_factor = ( 2 ** (len(self.vae.config.block_out_channels) - 1) if hasattr(self, "vae") and self.vae is not None else 8 ) diff --git a/vllm_omni/diffusion/models/z_image/z_image_transformer.py b/vllm_omni/diffusion/models/z_image/z_image_transformer.py index 0d7e68f70..cf04bf172 100644 --- a/vllm_omni/diffusion/models/z_image/z_image_transformer.py +++ b/vllm_omni/diffusion/models/z_image/z_image_transformer.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# _sp_plan definition adapted from HuggingFace diffusers library (_cp_plan) # Copyright 2025 Alibaba Z-Image Team and The HuggingFace Team. All rights reserved. # @@ -32,7 +33,14 @@ from vllm.model_executor.layers.linear import ( from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm_omni.diffusion.attention.layer import Attention -from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available +from vllm_omni.diffusion.distributed.sp_plan import ( + SequenceParallelInput, + SequenceParallelOutput, +) +from vllm_omni.diffusion.forward_context import ( + get_forward_context, + is_forward_context_available, +) from vllm_omni.diffusion.layers.rope import RotaryEmbedding ADALN_EMBED_DIM = 256 @@ -41,6 +49,63 @@ SEQ_MULTI_OF = 32 logger = init_logger(__name__) +class UnifiedPrepare(nn.Module): + """Prepares unified tensors for transformer blocks. + + This module encapsulates the unification of x and cap tensors into unified + sequences. Similar to how Wan's `rope` module outputs rotary embeddings, + this module outputs unified tensors that can be sharded via _sp_plan's + split_output=True mechanism. + + This follows the diffusers pattern where tensor preparation happens in + a dedicated submodule, enabling _sp_plan hooks to work at module boundaries. + """ + + def forward( + self, + x: torch.Tensor, + x_cos: torch.Tensor, + x_sin: torch.Tensor, + cap_feats: torch.Tensor, + cap_cos: torch.Tensor, + cap_sin: torch.Tensor, + x_item_seqlens: list[int], + cap_item_seqlens: list[int], + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Combine x and cap tensors into unified sequences. + + Returns: + unified: Combined hidden states [batch, seq_len, dim] + unified_cos: Combined RoPE cos [batch, seq_len, rope_dim] + unified_sin: Combined RoPE sin [batch, seq_len, rope_dim] + unified_attn_mask: Combined attention mask [batch, seq_len] + """ + bsz = x.shape[0] + device = x.device + + unified = [] + unified_cos = [] + unified_sin = [] + for i in range(bsz): + x_len = x_item_seqlens[i] + cap_len = cap_item_seqlens[i] + unified.append(torch.cat([x[i][:x_len], cap_feats[i][:cap_len]])) + unified_cos.append(torch.cat([x_cos[i][:x_len], cap_cos[i][:cap_len]])) + unified_sin.append(torch.cat([x_sin[i][:x_len], cap_sin[i][:cap_len]])) + + unified_item_seqlens = [a + b for a, b in zip(cap_item_seqlens, x_item_seqlens)] + unified_max_item_seqlen = max(unified_item_seqlens) + + unified = pad_sequence(unified, batch_first=True, padding_value=0.0) + unified_cos = pad_sequence(unified_cos, batch_first=True, padding_value=0.0) + unified_sin = pad_sequence(unified_sin, batch_first=True, padding_value=0.0) + unified_attn_mask = torch.zeros((bsz, unified_max_item_seqlen), dtype=torch.bool, device=device) + for i, seq_len in enumerate(unified_item_seqlens): + unified_attn_mask[i, :seq_len] = 1 + + return unified, unified_cos, unified_sin, unified_attn_mask + + def _positive_divisors(n: int) -> set[int]: if n <= 0: return set() @@ -456,8 +521,52 @@ class RopeEmbedder: class ZImageTransformer2DModel(nn.Module): + """Z-Image Transformer model for image generation. + + Sequence Parallelism: + This model supports non-intrusive SP via _sp_plan. The plan specifies: + - Input splitting at first main transformer block (unified sequence) + - RoPE (cos/sin) splitting along sequence dimension + - Attention mask splitting along sequence dimension + - Output gathering at final_layer + + The SP is applied to the main `layers` transformer blocks where the + unified image+caption sequence is processed jointly. + + Note: noise_refiner and context_refiner are NOT parallelized as they + process image and caption separately before unification. + + Important: The default _sp_plan assumes patch_size=2 and f_patch_size=1. + If using different patch configurations, update _sp_plan accordingly. + + Note: Our "Sequence Parallelism" (SP) corresponds to "Context Parallelism" (CP) in diffusers. + """ + _repeated_blocks = ["ZImageTransformerBlock"] + # Sequence Parallelism for Z-Image (following diffusers' _cp_plan pattern) + # Similar to how Wan uses `rope` module's split_output to shard rotary embeddings, + # Z-Image uses `unified_prepare` module's split_output to shard unified tensors. + # + # The _sp_plan specifies sharding/gathering at module boundaries: + # - unified_prepare: Split all 4 outputs (unified, cos, sin, attn_mask) via split_output=True + # - layers.0: hidden_states input is already sharded from unified_prepare output + # - all_final_layer.2-1: Gather outputs after the final layer + # + # Note: _sp_plan corresponds to diffusers' _cp_plan (Context Parallelism) + _sp_plan = { + # Shard unified_prepare outputs (similar to Wan's rope module) + # This shards all 4 return values: unified, unified_cos, unified_sin, unified_attn_mask + "unified_prepare": { + 0: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True), # unified + 1: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True), # unified_cos + 2: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True), # unified_sin + 3: SequenceParallelInput(split_dim=1, expected_dims=2, split_output=True), # unified_attn_mask + }, + # Gather output at final_layer (default: patch_size=2, f_patch_size=1) + "all_final_layer.2-1": SequenceParallelOutput(gather_dim=1, expected_dims=3), + } + def __init__( self, all_patch_size=(2,), @@ -572,6 +681,11 @@ class ZImageTransformer2DModel(nn.Module): self.rope_embedder = RopeEmbedder(theta=rope_theta, axes_dims=axes_dims, axes_lens=axes_lens) + # UnifiedPrepare module for combining x and cap tensors + # This enables _cp_plan to shard outputs via split_output=True + # Similar to how Wan's rope module enables rotary embedding sharding + self.unified_prepare = UnifiedPrepare() + def unpatchify(self, x: list[torch.Tensor], size: list[tuple], patch_size, f_patch_size) -> list[torch.Tensor]: pH = pW = patch_size pF = f_patch_size @@ -772,31 +886,19 @@ class ZImageTransformer2DModel(nn.Module): for layer in self.context_refiner: cap_feats = layer(cap_feats, cap_attn_mask, cap_cos, cap_sin) - # unified - unified = [] - unified_cos = [] - unified_sin = [] - for i in range(bsz): - x_len = x_item_seqlens[i] - cap_len = cap_item_seqlens[i] - unified.append(torch.cat([x[i][:x_len], cap_feats[i][:cap_len]])) - unified_cos.append(torch.cat([x_cos[i][:x_len], cap_cos[i][:cap_len]])) - unified_sin.append(torch.cat([x_sin[i][:x_len], cap_sin[i][:cap_len]])) - unified_item_seqlens = [a + b for a, b in zip(cap_item_seqlens, x_item_seqlens)] - assert unified_item_seqlens == [len(_) for _ in unified] - unified_max_item_seqlen = max(unified_item_seqlens) - - unified = pad_sequence(unified, batch_first=True, padding_value=0.0) - unified_cos = pad_sequence(unified_cos, batch_first=True, padding_value=0.0) - unified_sin = pad_sequence(unified_sin, batch_first=True, padding_value=0.0) - unified_attn_mask = torch.zeros((bsz, unified_max_item_seqlen), dtype=torch.bool, device=device) - for i, seq_len in enumerate(unified_item_seqlens): - unified_attn_mask[i, :seq_len] = 1 + # Prepare unified tensors via UnifiedPrepare module + # This enables _cp_plan to shard outputs via split_output=True + unified, unified_cos, unified_sin, unified_attn_mask = self.unified_prepare( + x, x_cos, x_sin, cap_feats, cap_cos, cap_sin, x_item_seqlens, cap_item_seqlens + ) + # Main transformer blocks for layer in self.layers: unified = layer(unified, unified_attn_mask, unified_cos, unified_sin, adaln_input) + # Final layer unified = self.all_final_layer[f"{patch_size}-{f_patch_size}"](unified, adaln_input) + unified = list(unified.unbind(dim=0)) x = self.unpatchify(unified, x_size, patch_size, f_patch_size) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 99b0ea4c7..466a58e2c 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -3,9 +3,15 @@ import importlib +import torch.nn as nn +from vllm.logger import init_logger from vllm.model_executor.models.registry import _LazyRegisteredModel, _ModelRegistry from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig, get_sp_plan_from_model +from vllm_omni.diffusion.hooks.sequence_parallel import apply_sequence_parallel + +logger = init_logger(__name__) _DIFFUSION_MODELS = { # arch:(mod_folder, mod_relname, cls_name) @@ -100,7 +106,24 @@ DiffusionModelRegistry = _ModelRegistry( def initialize_model( od_config: OmniDiffusionConfig, -): +) -> nn.Module: + """Initialize a diffusion model from the registry. + + This function: + 1. Loads the model class from the registry + 2. Instantiates the model with the config + 3. Configures VAE optimization settings + 4. Applies sequence parallelism if enabled (similar to diffusers' enable_parallelism) + + Args: + od_config: The OmniDiffusion configuration. + + Returns: + The initialized pipeline model. + + Raises: + ValueError: If the model class is not found in the registry. + """ model_class = DiffusionModelRegistry._try_load_model_cls(od_config.model_class_name) if model_class is not None: model = model_class(od_config=od_config) @@ -110,11 +133,73 @@ def initialize_model( if hasattr(model.vae, "use_tiling"): model.vae.use_tiling = od_config.vae_use_tiling + # Apply sequence parallelism if enabled + # This follows diffusers' pattern where enable_parallelism() is called + # at model loading time, not inside individual model files + _apply_sequence_parallel_if_enabled(model, od_config) + return model else: raise ValueError(f"Model class {od_config.model_class_name} not found in diffusion model registry.") +def _apply_sequence_parallel_if_enabled(model, od_config: OmniDiffusionConfig) -> None: + """Apply sequence parallelism hooks if SP is enabled. + + This is the centralized location for enabling SP, similar to diffusers' + ModelMixin.enable_parallelism() method. It applies _sp_plan hooks to + transformer models that define them. + + Note: Our "Sequence Parallelism" (SP) corresponds to "Context Parallelism" (CP) in diffusers. + We use _sp_plan instead of diffusers' _cp_plan. + + Args: + model: The pipeline model (e.g., ZImagePipeline). + od_config: The OmniDiffusion configuration. + """ + + try: + sp_size = od_config.parallel_config.sequence_parallel_size + if sp_size <= 1: + return + + # Find transformer model(s) in the pipeline that have _sp_plan + transformer_attrs = ["transformer", "dit", "unet"] + for attr in transformer_attrs: + if not hasattr(model, attr): + continue + + transformer = getattr(model, attr) + if transformer is None: + continue + + plan = get_sp_plan_from_model(transformer) + if plan is None: + continue + + # Create SP config + sp_config = SequenceParallelConfig( + ulysses_degree=od_config.parallel_config.ulysses_degree, + ring_degree=od_config.parallel_config.ring_degree, + ) + + # Apply hooks according to the plan + mode = ( + "hybrid" + if sp_config.ulysses_degree > 1 and sp_config.ring_degree > 1 + else ("ulysses" if sp_config.ulysses_degree > 1 else "ring") + ) + logger.info( + f"Applying sequence parallelism to {transformer.__class__.__name__} " + f"(sp_size={sp_size}, mode={mode}, ulysses={sp_config.ulysses_degree}, ring={sp_config.ring_degree})" + ) + apply_sequence_parallel(transformer, sp_config, plan) + return # Only apply to first transformer found + + except Exception as e: + logger.warning(f"Failed to apply sequence parallelism: {e}. Continuing without SP hooks.") + + _DIFFUSION_POST_PROCESS_FUNCS = { # arch: post_process_func # `post_process_func` function must be placed in {mod_folder}/{mod_relname}.py,