diff --git a/docs/user_guide/acceleration/parallelism_acceleration.md b/docs/user_guide/acceleration/parallelism_acceleration.md index 31852c6d2..cf18b18e5 100644 --- a/docs/user_guide/acceleration/parallelism_acceleration.md +++ b/docs/user_guide/acceleration/parallelism_acceleration.md @@ -6,7 +6,7 @@ This guide includes how to use parallelism methods in vLLM-Omni to speed up diff The following parallelism methods are currently supported in vLLM-Omni: -1. DeepSpeed Ulysses Sequence Parallel (Ulysses-SP) ([paper](https://arxiv.org/pdf/2309.14509)): Ulysses-SP splits the input along the sequence dimension and uses all-to-all communication to allow each device to compute only a subset of attention heads. +1. DeepSpeed Ulysses Sequence Parallel (DeepSpeed Ulysses-SP) ([arxiv paper](https://arxiv.org/pdf/2309.14509)): Ulysses-SP splits the input along the sequence dimension and uses all-to-all communication to allow each device to compute only a subset of attention heads. The following table shows which models are currently supported by parallelism method: @@ -22,9 +22,9 @@ The following table shows which models are currently supported by parallelism me #### Ulysses-SP -##### Quick Start +##### Offline Inference -An example of using Ulysses-SP is shown below: +An example of offline inference script using [Ulysses-SP](https://arxiv.org/pdf/2309.14509) is shown below: ```python from vllm_omni import Omni from vllm_omni.diffusion.data import DiffusionParallelConfig @@ -70,56 +70,84 @@ 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 instruction on how to parallelize this model with Ulysses-SP. +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). -First, please edit the `TransformerModel`'s `forward` function in the `xxx_model_transformer.py` to make the inputs (image hidden states, positional embeddings, etc.) as chunks separated at the sequence dimension. Taking `qwen_image_transformer.py` as an example: +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: + +- 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. + +First, add the sequence-parallel helpers and (for Qwen-Image) the forward-context flag. Then, in the transformer's `forward()`: + +- 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. + +Taking `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` as an example: ```diff -class QwenImageTransformer2DModel(nn.Module): +from vllm_omni.diffusion.distributed.parallel_state import ( + get_sequence_parallel_rank, + get_sequence_parallel_world_size, + get_sp_group, +) +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: -+ hidden_states = torch.chunk(hidden_states, get_sequence_parallel_world_size(), dim=-2)[ -+ get_sequence_parallel_rank() -+ ] ++ 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) + hidden_states = self.img_in(hidden_states) - ... - image_rotary_emb = self.pos_embed(img_shapes, txt_seq_lens, device=hidden_states.device) + ... + image_rotary_emb = self.pos_embed(img_shapes, txt_seq_lens, device=hidden_states.device) -+ def get_rotary_emb_chunk(freqs): -+ 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) -+ image_rotary_emb = (img_freqs, txt_freqs) ++ 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, please call `get_sp_group().all_gather` to gather the chunked outputs across devices, and concatenate them at the sequence dimension. - +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): ```diff -class QwenImageTransformer2DModel(nn.Module): - def forward( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor = None, +class QwenImageTransformer2DModel(...): + def forward(...): ... - ): - # Use only the image part (hidden_states) from the dual-stream blocks - hidden_states = self.norm_out(hidden_states, temb) - output = self.proj_out(hidden_states) + 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) - return Transformer2DModelOutput(sample=output) ++ if self.parallel_config.sequence_parallel_size > 1: ++ output = get_sp_group().all_gather(output, dim=-2) + return Transformer2DModelOutput(sample=output) ``` Finally, you can set the parallel configuration and pass it to `Omni` and start parallel inference with: diff --git a/tests/diffusion/attention/test_ulysses_sequence_parallel.py b/tests/diffusion/attention/test_ulysses_sequence_parallel.py index 33429142c..0b60a340c 100644 --- a/tests/diffusion/attention/test_ulysses_sequence_parallel.py +++ b/tests/diffusion/attention/test_ulysses_sequence_parallel.py @@ -8,6 +8,9 @@ import pytest import torch from vllm.platforms import current_platform +from vllm_omni.diffusion.attention.backends.abstract import ( + AttentionMetadata, +) from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.data import ( DiffusionParallelConfig, @@ -16,9 +19,11 @@ from vllm_omni.diffusion.data import ( ) from vllm_omni.diffusion.distributed.parallel_state import ( destroy_distributed_env, + get_sequence_parallel_world_size, init_distributed_environment, initialize_model_parallel, ) +from vllm_omni.diffusion.forward_context import set_forward_context from vllm_omni.utils.platform_utils import detect_device_type device_type = detect_device_type() @@ -29,6 +34,9 @@ elif device_type == "npu": else: raise ValueError(f"Unsupported device type: {device_type} for this test script! Expected GPU or NPU.") +global split_text_embed_in_sp +split_text_embed_in_sp = False + def update_environment_variables(envs_dict: dict[str, str]): """Update multiple environment variables with logging.""" @@ -70,26 +78,52 @@ class TestAttentionModel(torch.nn.Module): self.v_proj = torch.nn.Linear(hidden_size, (num_kv_heads or num_heads) * head_size) self.o_proj = torch.nn.Linear(num_heads * head_size, hidden_size) - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + def forward(self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | None = None) -> torch.Tensor: """Forward pass through attention layer.""" batch_size, seq_len, _ = hidden_states.shape + # Combine hidden_states and encoder_hidden_states if provided + if encoder_hidden_states is not None: + # Concatenate along sequence dimension + combined_hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) + else: + combined_hidden_states = hidden_states + # Project to Q, K, V - q = self.q_proj(hidden_states) - k = self.k_proj(hidden_states) - v = self.v_proj(hidden_states) + q = self.q_proj(combined_hidden_states) + k = self.k_proj(combined_hidden_states) + v = self.v_proj(combined_hidden_states) - # Reshape to (batch_size, seq_len, num_heads, head_size) - q = q.view(batch_size, seq_len, self.num_heads, self.head_size) - k = k.view(batch_size, seq_len, k.shape[-1] // self.head_size, self.head_size) - v = v.view(batch_size, seq_len, v.shape[-1] // self.head_size, self.head_size) + # Reshape to (batch_size, total_seq_len, num_heads, head_size) + total_seq_len = combined_hidden_states.shape[1] + q = q.view(batch_size, total_seq_len, self.num_heads, self.head_size) + k = k.view(batch_size, total_seq_len, k.shape[-1] // self.head_size, self.head_size) + v = v.view(batch_size, total_seq_len, v.shape[-1] // self.head_size, self.head_size) - # Apply attention - attn_output = self.attention(q, k, v) + # Apply attention with split logic + if get_sequence_parallel_world_size() > 1 and not split_text_embed_in_sp and encoder_hidden_states is not None: + q_encoder, q_hidden = torch.split(q, [total_seq_len - seq_len, seq_len], dim=1) + k_encoder, k_hidden = torch.split(k, [total_seq_len - seq_len, seq_len], dim=1) + v_encoder, v_hidden = torch.split(v, [total_seq_len - seq_len, seq_len], dim=1) + + # Use hidden_states part as main attention, encoder part as joint + attn_output = self.attention( + q_hidden, + k_hidden, + v_hidden, + AttentionMetadata( + joint_query=q_encoder, + joint_key=k_encoder, + joint_value=v_encoder, + ), + ) + else: + attn_output = self.attention(q, k, v) # Reshape back and project - attn_output = attn_output.view(batch_size, seq_len, -1) + attn_output = attn_output.view(batch_size, total_seq_len, -1) output = self.o_proj(attn_output) + output = output[:, encoder_hidden_states.shape[1] :, :] if encoder_hidden_states is not None else output return output @@ -127,10 +161,10 @@ class TestMultiLayerAttentionModel(torch.nn.Module): ] ) - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + def forward(self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | None = None) -> torch.Tensor: """Forward pass through multiple attention layers.""" for layer in self.layers: - hidden_states = hidden_states + layer(hidden_states) + hidden_states = hidden_states + layer(hidden_states, encoder_hidden_states) return hidden_states @@ -141,15 +175,16 @@ class TestMultiLayerAttentionModel(torch.nn.Module): TestMultiLayerAttentionModel, ], ) -@pytest.mark.parametrize("ulysses_degree", [2, 4]) +@pytest.mark.parametrize("ulysses_degree", [2]) @pytest.mark.parametrize("ring_degree", [1]) @pytest.mark.parametrize("batch_size", [2]) @pytest.mark.parametrize("seq_len", [16]) +@pytest.mark.parametrize("encoder_seq_len", [16, 13]) # Test both divisible and non-divisible @pytest.mark.parametrize("num_heads", [8]) @pytest.mark.parametrize("head_size", [8]) -@pytest.mark.parametrize("causal", [True]) +@pytest.mark.parametrize("causal", [False]) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -@pytest.mark.parametrize("use_sync", [True, False]) +@pytest.mark.parametrize("use_sync", [True]) @pytest.mark.parametrize("dynamic", [False]) @pytest.mark.parametrize("use_compile", [False]) def test_ulysses_attention( @@ -163,12 +198,18 @@ def test_ulysses_attention( use_compile: bool, batch_size: int, seq_len: int, + encoder_seq_len: int, num_heads: int, head_size: int, ): """Test Ulysses attention by comparing with and without SP enabled.""" sequence_parallel_size = ulysses_degree * ring_degree + # Determine if we can split encoder_hidden_states in SP + can_split_encoder = (encoder_seq_len % sequence_parallel_size) == 0 + print(f"\nEncoder sequence length: {encoder_seq_len}, SP size: {sequence_parallel_size}") + print(f"Can split encoder in SP: {can_split_encoder}") + # Create temporary files to share results between processes with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as f: baseline_output_file = f.name @@ -178,6 +219,8 @@ def test_ulysses_attention( model_state_file = f.name with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as f: input_data_file = f.name + with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as f: + encoder_input_data_file = f.name try: # Step 1: Run without SP (baseline with ulysses_degree=1, ring_degree=1) @@ -189,6 +232,7 @@ def test_ulysses_attention( test_model_cls, batch_size, seq_len, + encoder_seq_len, num_heads, head_size, dtype, @@ -202,6 +246,7 @@ def test_ulysses_attention( baseline_output_file, model_state_file, input_data_file, + encoder_input_data_file, True, # is_baseline ), nprocs=1, @@ -216,6 +261,7 @@ def test_ulysses_attention( test_model_cls, batch_size, seq_len, + encoder_seq_len, num_heads, head_size, dtype, @@ -229,6 +275,7 @@ def test_ulysses_attention( sp_output_file, model_state_file, input_data_file, + encoder_input_data_file, False, # is_baseline ), nprocs=sequence_parallel_size, @@ -299,7 +346,7 @@ def test_ulysses_attention( finally: # Clean up temporary files - for f in [baseline_output_file, sp_output_file, model_state_file, input_data_file]: + for f in [baseline_output_file, sp_output_file, model_state_file, input_data_file, encoder_input_data_file]: if os.path.exists(f): os.remove(f) @@ -310,6 +357,7 @@ def ulysses_attention_on_test_model( test_model_cls: type[torch.nn.Module], batch_size: int, seq_len: int, + encoder_seq_len: int, num_heads: int, head_size: int, dtype: torch.dtype, @@ -323,6 +371,7 @@ def ulysses_attention_on_test_model( output_file: str, model_state_file: str, input_data_file: str, + encoder_input_data_file: str, is_baseline: bool, ): """Run Ulysses attention test on a test model and save results for comparison.""" @@ -330,8 +379,15 @@ def ulysses_attention_on_test_model( RANDOM_SEED = 42 current_platform.seed_everything(RANDOM_SEED) + # Determine if we can split encoder_hidden_states based on divisibility + global split_text_embed_in_sp + split_text_embed_in_sp = (encoder_seq_len % sequence_parallel_size) == 0 and not is_baseline + mode_str = "Baseline (no SP)" if is_baseline else f"SP (ulysses={ulysses_degree}, ring={ring_degree})" print(f"\n[{mode_str}] Rank {local_rank}/{world_size} - Random seed set to {RANDOM_SEED}") + print( + f"[{mode_str}] Rank {local_rank}/{world_size} - encoder_seq_len={encoder_seq_len}, split_text_embed_in_sp={split_text_embed_in_sp}" + ) device = torch.device(f"{device_type}:{local_rank}") torch_device.set_device(device) @@ -379,7 +435,7 @@ def ulysses_attention_on_test_model( ) # Set the config so Attention can access it - with set_current_omni_diffusion_config(od_config): + with set_forward_context(omni_diffusion_config=od_config), set_current_omni_diffusion_config(od_config): # Create model hidden_size = num_heads * head_size @@ -418,10 +474,17 @@ def ulysses_attention_on_test_model( dtype=dtype, device="cpu", ) + full_encoder_hidden_states = torch.randn( + (batch_size, encoder_seq_len, hidden_size), + dtype=dtype, + device="cpu", + ) with open(input_data_file, "wb") as f: pickle.dump(full_hidden_states.detach().cpu().float().numpy(), f) + with open(encoder_input_data_file, "wb") as f: + pickle.dump(full_encoder_hidden_states.detach().cpu().float().numpy(), f) - print("[Baseline] Saved model state and input data") + print("[Baseline] Saved model state and input data (including encoder_hidden_states)") # Synchronize to ensure baseline has saved data before SP loads it if world_size > 1: @@ -437,7 +500,13 @@ def ulysses_attention_on_test_model( full_hidden_states_np = pickle.load(f) full_hidden_states = torch.from_numpy(full_hidden_states_np).to(device).to(dtype) - print(f"[Rank {local_rank}] Loaded model state and full input data with shape {full_hidden_states.shape}") + with open(encoder_input_data_file, "rb") as f: + full_encoder_hidden_states_np = pickle.load(f) + full_encoder_hidden_states = torch.from_numpy(full_encoder_hidden_states_np).to(device).to(dtype) + + print( + f"[Rank {local_rank}] Loaded model state and full input data with shape {full_hidden_states.shape}, encoder shape {full_encoder_hidden_states.shape}" + ) # Split input sequence according to sequence parallel BEFORE model forward # Each rank gets a contiguous chunk of the sequence dimension @@ -446,10 +515,28 @@ def ulysses_attention_on_test_model( end_idx = start_idx + local_seq_len hidden_states = full_hidden_states[:, start_idx:end_idx, :].contiguous() - print( - f"[Rank {local_rank}] Split input: local_seq_len={local_seq_len}, " - f"indices=[{start_idx}:{end_idx}], local_shape={hidden_states.shape}" - ) + # Handle encoder_hidden_states splitting based on split_text_embed_in_sp + if get_sequence_parallel_world_size() > 1 and split_text_embed_in_sp: + # Split encoder_hidden_states in the same way as hidden_states + local_encoder_seq_len = encoder_seq_len // sequence_parallel_size + encoder_start_idx = local_rank * local_encoder_seq_len + encoder_end_idx = encoder_start_idx + local_encoder_seq_len + encoder_hidden_states = full_encoder_hidden_states[:, encoder_start_idx:encoder_end_idx, :].contiguous() + print( + f"[Rank {local_rank}] Split input: local_seq_len={local_seq_len}, " + f"indices=[{start_idx}:{end_idx}], hidden_states shape={hidden_states.shape}, " + f"encoder_hidden_states (split) shape={encoder_hidden_states.shape}, " + f"encoder_indices=[{encoder_start_idx}:{encoder_end_idx}]" + ) + else: + # No splitting for encoder_hidden_states, use full sequence + encoder_hidden_states = full_encoder_hidden_states + print( + f"[Rank {local_rank}] Split input: local_seq_len={local_seq_len}, " + f"indices=[{start_idx}:{end_idx}], hidden_states shape={hidden_states.shape}, " + f"encoder_hidden_states (full) shape={encoder_hidden_states.shape}, " + f"split_text_embed_in_sp={split_text_embed_in_sp}" + ) if dynamic: torch._dynamo.mark_dynamic(hidden_states, 0) @@ -461,26 +548,13 @@ def ulysses_attention_on_test_model( # Run forward pass with local sequence chunk print(f"[Rank {local_rank}] Running forward pass...") - output = model(hidden_states) + output = model(hidden_states, encoder_hidden_states) print(f"[Rank {local_rank}] Forward pass completed, output shape: {output.shape}") - - # Verify output shape assert output.shape == (batch_size, local_seq_len, hidden_size), ( f"Output shape mismatch: expected {(batch_size, local_seq_len, hidden_size)}, got {output.shape}" ) - # Verify SP usage for non-baseline runs - if not is_baseline: - if hasattr(model, "attention"): - assert hasattr(model.attention, "use_ulysses"), "Attention should have use_ulysses attribute" - assert model.attention.use_ulysses, "Attention should be using Ulysses" - elif hasattr(model, "layers"): - for i, layer in enumerate(model.layers): - assert hasattr(layer.attention, "use_ulysses"), ( - f"Layer {i} attention should have use_ulysses attribute" - ) - assert layer.attention.use_ulysses, f"Layer {i} attention should be using Ulysses" - + output = output.contiguous() # Gather outputs from all ranks AFTER computation if world_size > 1: print(f"[Rank {local_rank}] Gathering outputs from all {world_size} ranks...") @@ -491,7 +565,7 @@ def ulysses_attention_on_test_model( # Concatenate along sequence dimension to reconstruct full sequence full_output = torch.cat(gathered_outputs, dim=1) print(f"[Rank 0] Gathered and concatenated outputs: {full_output.shape}") - # Verify the full output shape matches expected + assert full_output.shape == (batch_size, seq_len, hidden_size), ( f"Gathered output shape mismatch: expected {(batch_size, seq_len, hidden_size)}, " f"got {full_output.shape}" @@ -512,9 +586,10 @@ def ulysses_attention_on_test_model( mode_str = "baseline (no SP)" if is_baseline else f"SP (ulysses={ulysses_degree}, ring={ring_degree})" print( f"\n[{mode_str}] ✓ Saved output with shape {full_output.shape}:\n" - f" - batch_size={batch_size}, seq_len={seq_len}\n" + f" - batch_size={batch_size}, seq_len={seq_len}, encoder_seq_len={encoder_seq_len}\n" f" - num_heads={num_heads}, head_size={head_size}\n" f" - dtype={dtype}, causal={causal}, use_sync={use_sync}\n" + f" - split_text_embed_in_sp={split_text_embed_in_sp}\n" ) destroy_distributed_env() diff --git a/tests/e2e/offline_inference/test_sequence_parallel.py b/tests/e2e/offline_inference/test_sequence_parallel.py index 7557162e3..e2d732b65 100644 --- a/tests/e2e/offline_inference/test_sequence_parallel.py +++ b/tests/e2e/offline_inference/test_sequence_parallel.py @@ -13,8 +13,10 @@ import os import sys from pathlib import Path +import numpy as np import pytest import torch +from PIL import Image # ruff: noqa: E402 REPO_ROOT = Path(__file__).resolve().parents[2] @@ -31,42 +33,113 @@ os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1" # Use random weights model for testing models = ["riverclouds/qwen_image_random"] +PROMPT = "a photo of a cat sitting on a laptop keyboard" + + +def _pil_to_float_rgb_tensor(img: Image.Image) -> torch.Tensor: + """Convert PIL image to float32 RGB tensor in [0, 1] with shape [H, W, 3].""" + arr = np.asarray(img.convert("RGB"), dtype=np.float32) / 255.0 + return torch.from_numpy(arr) + + +def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]: + """Return (mean_abs_diff, max_abs_diff) over RGB pixels in [0, 1].""" + ta = _pil_to_float_rgb_tensor(a) + tb = _pil_to_float_rgb_tensor(b) + assert ta.shape == tb.shape, f"Image shapes differ: {ta.shape} vs {tb.shape}" + abs_diff = torch.abs(ta - tb) + return abs_diff.mean().item(), abs_diff.max().item() + @pytest.mark.parametrize("model_name", models) -@pytest.mark.parametrize("ulysses_degree", [1, 2]) +@pytest.mark.parametrize("ulysses_degree", [2]) @pytest.mark.parametrize("ring_degree", [1]) -def test_sequence_parallel(model_name: str, ulysses_degree: int, ring_degree: int): - """Test SP (Ulysses-SP + Ring-SP) backend with diffusion model.""" - # Skip if not enough GPUs available +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_sequence_parallel(model_name: str, ulysses_degree: int, ring_degree: int, dtype: torch.dtype): + """Compare baseline (ulysses_degree=1) vs SP (ulysses_degree>1) outputs.""" + if ulysses_degree <= 1: + pytest.skip("This test compares ulysses_degree=1 vs ulysses_degree>1; provide ulysses_degree>1.") + + # Skip if not enough GPUs available for SP run if device_count() < ulysses_degree: pytest.skip(f"Test requires {ulysses_degree} GPUs but only {device_count()} available") - # Configure sequence parallel with DiffusionParallelConfig - parallel_config = DiffusionParallelConfig(ulysses_degree=ulysses_degree, ring_degree=ring_degree) - - m = Omni( - model=model_name, - parallel_config=parallel_config, - ) - # Use minimal settings for fast testing height = 256 width = 256 num_inference_steps = 4 # Minimal steps for fast test + seed = 42 - images = m.generate( - "a photo of a cat sitting on a laptop keyboard", - height=height, - width=width, - num_inference_steps=num_inference_steps, - guidance_scale=0.0, - generator=torch.Generator(get_device_name()).manual_seed(42), - num_outputs_per_prompt=1, # Single output for speed + # 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, + ) + try: + baseline_images = 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, + ) + finally: + baseline.close() + + assert baseline_images is not None + assert len(baseline_images) == 1 + assert baseline_images[0].width == width + assert baseline_images[0].height == height + + # 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, + ) + try: + sp_images = 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, + ) + finally: + sp.close() + + 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=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}" ) - # Verify generation succeeded - assert images is not None - assert len(images) == 1 - # Check image size - assert images[0].width == width - assert images[0].height == height + 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}" + ) diff --git a/vllm_omni/diffusion/attention/backends/abstract.py b/vllm_omni/diffusion/attention/backends/abstract.py index 5c9c78774..d574979d1 100644 --- a/vllm_omni/diffusion/attention/backends/abstract.py +++ b/vllm_omni/diffusion/attention/backends/abstract.py @@ -48,6 +48,14 @@ class AttentionBackend(ABC): @dataclass class AttentionMetadata: attn_mask: torch.Tensor | None = None + joint_query: torch.Tensor | None = None + # a replicated tensor among processes appended to the front or rear of query, depends the joint_strategy + joint_key: torch.Tensor | None = None + # a replicated tensor among processes appended to the front or rear of key, depends the joint_strategy + joint_value: torch.Tensor | None = None + # a replicated tensor among processes appended to the front or rear of value, depends the joint_strategy + joint_strategy: str = "front" + # the strategy to joint the query, key, and value, can be "front" or "rear" T = TypeVar("T", bound=AttentionMetadata) diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py index 99d4009c9..a565e6aef 100644 --- a/vllm_omni/diffusion/attention/layer.py +++ b/vllm_omni/diffusion/attention/layer.py @@ -7,17 +7,11 @@ # https://github.com/feifeibear/long-context-attention/blob/main/yunchang/attention/layer.py import torch -import torch.distributed as dist import torch.nn as nn -from torch import Tensor -from vllm_omni.diffusion.attention.backends.abstract import ( - AttentionMetadata, -) +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.parallel import build_parallel_attention_strategy from vllm_omni.diffusion.attention.selector import get_attn_backend -from vllm_omni.diffusion.data import get_current_omni_diffusion_config -from vllm_omni.diffusion.distributed.comm import SeqAllToAll4D -from vllm_omni.diffusion.distributed.parallel_state import get_sequence_parallel_world_size, get_sp_group class Attention(nn.Module): @@ -49,25 +43,13 @@ class Attention(nn.Module): self.scatter_idx = scatter_idx self.gather_idx = gather_idx self.use_sync = use_sync - self.ring_pg: dist.ProcessGroup | None = None - self.ulysses_pg: dist.ProcessGroup | None = None - self.use_ulysses = False - - try: - config = get_current_omni_diffusion_config() - if config.parallel_config.ulysses_degree > 1: - self.use_ulysses = True - # Get sequence parallel process group - try: - sp_group = get_sp_group() - self.ring_pg = sp_group.ring_group - self.ulysses_pg = sp_group.ulysses_group - assert get_sequence_parallel_world_size() > 1, "Sequence parallel world size must be > 1" - except (AssertionError, RuntimeError): - # If sequence parallel group is not initialized, disable Ulysses - self.use_ulysses = False - except Exception: - self.use_ulysses = False + # Parallel attention (communication / resharding) is a pluggable strategy. + # This keeps the attention kernel backend selection orthogonal. + self.parallel = build_parallel_attention_strategy( + scatter_idx=scatter_idx, + gather_idx=gather_idx, + use_sync=use_sync, + ) def forward( self, @@ -76,43 +58,13 @@ class Attention(nn.Module): value: torch.Tensor, attn_metadata: AttentionMetadata = None, ) -> torch.Tensor: - if self.use_ulysses: - return self._forward_ulysses(query, key, value, attn_metadata) - else: - # shape: (batch_size, seq_len, num_heads, head_size) - attn_output = self.attention.forward(query, key, value, attn_metadata) - return attn_output + # Parallel strategy may reshard/communicate QKV before the attention kernel. + query, key, value, attn_metadata, ctx = self.parallel.pre_attention(query, key, value, attn_metadata) - def _forward_ulysses( - self, - query: Tensor, - key: Tensor, - value: Tensor, - attn_metadata: AttentionMetadata = None, - ) -> Tensor: - """Ulysses attention forward pass with sequence parallelism.""" - # scatter 2, gather 1 - # (bs, seq_len/N, head_cnt, head_size) -> (bs, seq_len, head_cnt/N, head_size) - q = SeqAllToAll4D.apply(self.ulysses_pg, query, self.scatter_idx, self.gather_idx, self.use_sync) - k = SeqAllToAll4D.apply(self.ulysses_pg, key, self.scatter_idx, self.gather_idx, self.use_sync) - v = SeqAllToAll4D.apply(self.ulysses_pg, value, self.scatter_idx, self.gather_idx, self.use_sync) + # shape: (batch_size, seq_len, num_heads, head_size) + attn_output = self.attention.forward(query, key, value, attn_metadata) + if isinstance(attn_output, tuple): + attn_output = attn_output[0] - softmax_scale = self.softmax_scale - if softmax_scale is None: - softmax_scale = q.shape[-1] ** -0.5 - - context_layer = self.attention.forward( - q, - k, - v, - attn_metadata=attn_metadata, - ) - - if isinstance(context_layer, tuple): - context_layer = context_layer[0] - - # (bs, seq_len, head_cnt/N, head_size) -> (bs, seq_len/N, head_cnt, head_size) - # scatter 1, gather 2 - output = SeqAllToAll4D.apply(self.ulysses_pg, context_layer, self.gather_idx, self.scatter_idx, self.use_sync) - - return output + # Parallel strategy may need to reverse resharding after the kernel. + return self.parallel.post_attention(attn_output, ctx) diff --git a/vllm_omni/diffusion/attention/parallel/__init__.py b/vllm_omni/diffusion/attention/parallel/__init__.py new file mode 100644 index 000000000..49d776f0c --- /dev/null +++ b/vllm_omni/diffusion/attention/parallel/__init__.py @@ -0,0 +1,22 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Parallel attention strategies. + +This package provides **communication / resharding strategies** for attention, +orthogonal to the **attention kernel backend** (SDPA/Flash/Sage). + +The goal is to keep `vllm_omni.diffusion.attention.layer.Attention` small and +extensible: adding a new parallelism method should not require editing the core +Attention module, only adding a new strategy and selecting it in the factory. +""" + +from .base import NoParallelAttention, ParallelAttentionContext, ParallelAttentionStrategy +from .factory import build_parallel_attention_strategy + +__all__ = [ + "ParallelAttentionStrategy", + "ParallelAttentionContext", + "NoParallelAttention", + "build_parallel_attention_strategy", +] diff --git a/vllm_omni/diffusion/attention/parallel/base.py b/vllm_omni/diffusion/attention/parallel/base.py new file mode 100644 index 000000000..c98add422 --- /dev/null +++ b/vllm_omni/diffusion/attention/parallel/base.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Protocol + +import torch + +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata + + +@dataclass(frozen=True, slots=True) +class ParallelAttentionContext: + """Opaque per-forward context returned by a parallel strategy. + + Strategies may stash whatever they need here to finish post-processing after + the attention kernel runs (e.g. reverse resharding, slicing metadata, etc.). + """ + + name: str + + +class ParallelAttentionStrategy(Protocol): + """Pluggable strategy for parallel attention communication/resharding. + + This is intentionally orthogonal to the attention *kernel* backend. + The kernel backend implements `AttentionImpl.forward()` for a given device, + while the parallel strategy implements how Q/K/V and outputs are sharded / + communicated across ranks. + """ + + @property + def enabled(self) -> bool: ... + + @property + def name(self) -> str: ... + + def pre_attention( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata | None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, AttentionMetadata | None, ParallelAttentionContext | None]: + """Runs before the attention kernel. + + Returns possibly transformed Q/K/V and metadata, and an optional context + for `post_attention`. + """ + + def post_attention( + self, + attn_output: torch.Tensor, + ctx: ParallelAttentionContext | None, + ) -> torch.Tensor: + """Runs after the attention kernel.""" + + +class NoParallelAttention: + """Default strategy: do nothing (single device / no SP).""" + + @property + def enabled(self) -> bool: + return False + + @property + def name(self) -> str: + return "none" + + def pre_attention( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata | None, + ): + return query, key, value, attn_metadata, None + + def post_attention(self, attn_output: torch.Tensor, ctx: ParallelAttentionContext | None) -> torch.Tensor: + return attn_output diff --git a/vllm_omni/diffusion/attention/parallel/factory.py b/vllm_omni/diffusion/attention/parallel/factory.py new file mode 100644 index 000000000..ee0113fc9 --- /dev/null +++ b/vllm_omni/diffusion/attention/parallel/factory.py @@ -0,0 +1,45 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention, ParallelAttentionStrategy +from vllm_omni.diffusion.attention.parallel.ulysses import UlyssesParallelAttention +from vllm_omni.diffusion.data import get_current_omni_diffusion_config +from vllm_omni.diffusion.distributed.parallel_state import get_sp_group + + +def build_parallel_attention_strategy( + *, + scatter_idx: int, + gather_idx: int, + use_sync: bool, +) -> ParallelAttentionStrategy: + """Select a parallel attention strategy based on current diffusion config. + + Design principle: + - Attention kernel backend selection remains in `attention/selector.py`. + - Parallel attention selection is handled here, based on distributed config + and initialized process groups. + """ + try: + cfg = get_current_omni_diffusion_config() + p = cfg.parallel_config + except Exception: + return NoParallelAttention() + + # Current implementation supports Ulysses sequence-parallel attention. + # We intentionally do NOT infer ring attention from ring_degree here yet. + if getattr(p, "ulysses_degree", 1) > 1: + try: + sp_group = get_sp_group() + except Exception: + return NoParallelAttention() + return UlyssesParallelAttention( + sp_group=sp_group, + scatter_idx=scatter_idx, + gather_idx=gather_idx, + use_sync=use_sync, + ) + + return NoParallelAttention() diff --git a/vllm_omni/diffusion/attention/parallel/ulysses.py b/vllm_omni/diffusion/attention/parallel/ulysses.py new file mode 100644 index 000000000..acf468486 --- /dev/null +++ b/vllm_omni/diffusion/attention/parallel/ulysses.py @@ -0,0 +1,135 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from dataclasses import dataclass + +import torch +import torch.distributed as dist + +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.parallel.base import ParallelAttentionContext +from vllm_omni.diffusion.distributed.comm import SeqAllToAll4D +from vllm_omni.diffusion.distributed.group_coordinator import SequenceParallelGroupCoordinator + + +@dataclass(frozen=True, slots=True) +class _UlyssesCtx(ParallelAttentionContext): + """Per-forward context for Ulysses sequence-parallel attention.""" + + ulysses_pg: dist.ProcessGroup + scatter_idx: int + gather_idx: int + use_sync: bool + + +class UlyssesParallelAttention: + """Ulysses sequence-parallel strategy (all-to-all over seq/head dims). + + This preserves the semantics previously implemented in + `Attention._forward_ulysses`: + - If `AttentionMetadata.joint_*` is provided, joint_query is concatenated to + query *before* all-to-all; joint_key/value are concatenated *after* all-to-all. + - joint_key/value are assumed to be replicated across SP ranks and are sliced + by ulysses head rank before concatenation. + """ + + def __init__( + self, + sp_group: SequenceParallelGroupCoordinator, + scatter_idx: int, + gather_idx: int, + use_sync: bool, + ) -> None: + self._sp_group = sp_group + self._ulysses_pg = sp_group.ulysses_group + self._scatter_idx = scatter_idx + self._gather_idx = gather_idx + self._use_sync = use_sync + + @property + def enabled(self) -> bool: + return True + + @property + def name(self) -> str: + return "ulysses" + + def pre_attention( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata | None, + ): + joint_tensor_query = joint_tensor_key = joint_tensor_value = None + joint_strategy = None + + if attn_metadata is not None: + joint_tensor_query = attn_metadata.joint_query + joint_tensor_key = attn_metadata.joint_key + joint_tensor_value = attn_metadata.joint_value + joint_strategy = attn_metadata.joint_strategy + + is_joint = False + if joint_tensor_query is not None and joint_tensor_key is not None and joint_tensor_value is not None: + supported_joint_strategy = ["front", "rear"] + if joint_strategy not in supported_joint_strategy: + raise ValueError( + f"joint_strategy: {joint_strategy} not supported." + f" supported joint strategy: {supported_joint_strategy}" + ) + if joint_strategy == "rear": + query = torch.cat([query, joint_tensor_query], dim=1) + else: + query = torch.cat([joint_tensor_query, query], dim=1) + is_joint = True + elif joint_tensor_query is None and joint_tensor_key is None and joint_tensor_value is None: + pass + else: + raise ValueError("joint_query, joint_key, and joint_value should be None or not None simultaneously.") + + if is_joint: + # Slice joint key/value heads for this ulysses rank. + ulysses_world_size = self._sp_group.ulysses_world_size + ulysses_rank = self._sp_group.ulysses_rank + attn_heads_per_ulysses_rank = joint_tensor_key.shape[-2] // ulysses_world_size + joint_tensor_key = joint_tensor_key[ + ..., + attn_heads_per_ulysses_rank * ulysses_rank : attn_heads_per_ulysses_rank * (ulysses_rank + 1), + :, + ] + joint_tensor_value = joint_tensor_value[ + ..., + attn_heads_per_ulysses_rank * ulysses_rank : attn_heads_per_ulysses_rank * (ulysses_rank + 1), + :, + ] + + # (bs, seq_len/P, head_cnt, head_size) -> (bs, seq_len, head_cnt/P, head_size) + query = SeqAllToAll4D.apply(self._ulysses_pg, query, self._scatter_idx, self._gather_idx, self._use_sync) + key = SeqAllToAll4D.apply(self._ulysses_pg, key, self._scatter_idx, self._gather_idx, self._use_sync) + value = SeqAllToAll4D.apply(self._ulysses_pg, value, self._scatter_idx, self._gather_idx, self._use_sync) + + if is_joint: + # Concatenate joint key/value after all-to-all (matches previous implementation). + if joint_strategy == "front": + key = torch.cat([joint_tensor_key, key], dim=1) + value = torch.cat([joint_tensor_value, value], dim=1) + else: # "rear" + key = torch.cat([key, joint_tensor_key], dim=1) + value = torch.cat([value, joint_tensor_value], dim=1) + + ctx = _UlyssesCtx( + name=self.name, + ulysses_pg=self._ulysses_pg, + scatter_idx=self._scatter_idx, + gather_idx=self._gather_idx, + use_sync=self._use_sync, + ) + return query, key, value, attn_metadata, ctx + + def post_attention(self, attn_output: torch.Tensor, ctx: ParallelAttentionContext | None) -> torch.Tensor: + assert isinstance(ctx, _UlyssesCtx), f"Unexpected ctx type: {type(ctx)!r}" + # Reverse: (bs, seq_len, head_cnt/P, head_size) -> (bs, seq_len/P, head_cnt, head_size) + return SeqAllToAll4D.apply(ctx.ulysses_pg, attn_output, ctx.gather_idx, ctx.scatter_idx, ctx.use_sync) diff --git a/vllm_omni/diffusion/forward_context.py b/vllm_omni/diffusion/forward_context.py new file mode 100644 index 000000000..3c33f8105 --- /dev/null +++ b/vllm_omni/diffusion/forward_context.py @@ -0,0 +1,90 @@ +from contextlib import contextmanager +from dataclasses import dataclass + +from vllm.config import VllmConfig + +from vllm_omni.diffusion.attention.backends.abstract import ( + AttentionMetadata, +) +from vllm_omni.diffusion.data import OmniDiffusionConfig + + +@dataclass +class ForwardContext: + """ + set forward context for diffusion models + """ + + vllm_config: VllmConfig | None = None + omni_diffusion_config: OmniDiffusionConfig | None = None + attn_metadata: dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None = None + 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 + + def __post_init__(self): + pass + + +_forward_context: ForwardContext | None = None + + +def get_forward_context() -> ForwardContext: + """Get the current forward context.""" + assert _forward_context is not None, ( + "Forward context is not set. Please use `set_forward_context` to set the forward context." + ) + return _forward_context + + +def is_forward_context_available() -> bool: + return _forward_context is not None + + +def create_forward_context( + vllm_config: VllmConfig | None = None, + omni_diffusion_config: OmniDiffusionConfig | None = None, + attn_metadata: dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None = None, + split_text_embed_in_sp: bool = False, +): + return ForwardContext( + vllm_config=vllm_config, + omni_diffusion_config=omni_diffusion_config, + attn_metadata=attn_metadata, + split_text_embed_in_sp=split_text_embed_in_sp, + ) + + +@contextmanager +def override_forward_context(forward_context: ForwardContext | None): + """A context manager that overrides the current forward context. + This is used to override the forward context for a specific + forward pass. + """ + global _forward_context + prev_context = _forward_context + _forward_context = forward_context + try: + yield + finally: + _forward_context = prev_context + + +@contextmanager +def set_forward_context( + vllm_config: VllmConfig | None = None, + omni_diffusion_config: OmniDiffusionConfig | None = None, + attn_metadata: dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None = None, + split_text_embed_in_sp: bool = False, +): + """A context manager that stores the current forward context, + can be attention metadata, split_text_embed_in_sp, etc. + Here we can inject common logic for every model forward pass. + """ + forward_context = create_forward_context( + vllm_config=vllm_config, + omni_diffusion_config=omni_diffusion_config, + attn_metadata=attn_metadata, + split_text_embed_in_sp=split_text_embed_in_sp, + ) + with override_forward_context(forward_context): + yield 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 5b961d617..e6d10b087 100644 --- a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py +++ b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py @@ -19,14 +19,18 @@ from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import QKVParallelLinear, ReplicatedLinear from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm_omni.diffusion.attention.backends.abstract import ( + AttentionMetadata, +) from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.cache.base import CachedTransformer -from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.data import OmniDiffusionConfig, get_current_omni_diffusion_config from vllm_omni.diffusion.distributed.parallel_state import ( get_sequence_parallel_rank, get_sequence_parallel_world_size, get_sp_group, ) +from vllm_omni.diffusion.forward_context import get_forward_context from vllm_omni.diffusion.layers.rope import RotaryEmbedding logger = init_logger(__name__) @@ -356,6 +360,12 @@ class QwenImageCrossAttention(nn.Module): ) self.rope = RotaryEmbedding(is_neox_style=False) + try: + config = get_current_omni_diffusion_config() + self.parallel_config = config.parallel_config + except Exception: + self.parallel_config = None + def forward( self, hidden_states: torch.Tensor, @@ -405,11 +415,31 @@ class QwenImageCrossAttention(nn.Module): joint_value = torch.cat([txt_value, img_value], dim=1) # Compute joint attention - joint_hidden_states = self.attn( - joint_query, - joint_key, - joint_value, - ) + + if ( + self.parallel_config is not None + and self.parallel_config.sequence_parallel_size > 1 + and not get_forward_context().split_text_embed_in_sp + ): + # if using sequence parallel, but not splitting text embed, + # we need to pass text embedding to attention layer as joint qkv + joint_hidden_states = self.attn( + img_query, + img_key, + img_value, + AttentionMetadata( + joint_query=txt_query, + joint_key=txt_key, + joint_value=txt_value, + joint_strategy="front", + ), + ) + else: + joint_hidden_states = self.attn( + joint_query, + joint_key, + joint_value, + ) joint_hidden_states = joint_hidden_states.flatten(2, 3) joint_hidden_states = joint_hidden_states.to(joint_query.dtype) @@ -709,7 +739,10 @@ class QwenImageTransformer2DModel(CachedTransformer): hidden_states = torch.chunk(hidden_states, get_sequence_parallel_world_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) # Ensure timestep tensor is on the same device and dtype as hidden_states @@ -746,6 +779,8 @@ class QwenImageTransformer2DModel(CachedTransformer): 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) for index_block, block in enumerate(self.transformer_blocks): diff --git a/vllm_omni/diffusion/worker/gpu_worker.py b/vllm_omni/diffusion/worker/gpu_worker.py index e7fd05f70..dfaafa5de 100644 --- a/vllm_omni/diffusion/worker/gpu_worker.py +++ b/vllm_omni/diffusion/worker/gpu_worker.py @@ -23,6 +23,7 @@ from vllm_omni.diffusion.distributed.parallel_state import ( init_distributed_environment, initialize_model_parallel, ) +from vllm_omni.diffusion.forward_context import set_forward_context from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.request import OmniDiffusionRequest @@ -65,31 +66,33 @@ class GPUWorker: vllm_config = VllmConfig() vllm_config.parallel_config.tensor_parallel_size = self.od_config.parallel_config.tensor_parallel_size vllm_config.parallel_config.data_parallel_size = self.od_config.parallel_config.data_parallel_size + self.vllm_config = vllm_config + with ( + set_current_omni_diffusion_config(self.od_config), + set_current_vllm_config(vllm_config), + ): + init_distributed_environment(world_size=world_size, rank=rank) + logger.info(f"Worker {self.rank}: Initialized device and distributed environment.") + parallel_config = self.od_config.parallel_config + initialize_model_parallel( + data_parallel_size=parallel_config.data_parallel_size, + cfg_parallel_size=parallel_config.cfg_parallel_size, + sequence_parallel_size=parallel_config.sequence_parallel_size, + ulysses_degree=parallel_config.ulysses_degree, + ring_degree=parallel_config.ring_degree, + tensor_parallel_size=parallel_config.tensor_parallel_size, + pipeline_parallel_size=parallel_config.pipeline_parallel_size, + ) - with set_current_omni_diffusion_config(self.od_config): - with set_current_vllm_config(vllm_config): - init_distributed_environment(world_size=world_size, rank=rank) - logger.info(f"Worker {self.rank}: Initialized device and distributed environment.") - parallel_config = self.od_config.parallel_config - initialize_model_parallel( - data_parallel_size=parallel_config.data_parallel_size, - cfg_parallel_size=parallel_config.cfg_parallel_size, - sequence_parallel_size=parallel_config.sequence_parallel_size, - ulysses_degree=parallel_config.ulysses_degree, - ring_degree=parallel_config.ring_degree, - tensor_parallel_size=parallel_config.tensor_parallel_size, - pipeline_parallel_size=parallel_config.pipeline_parallel_size, + load_config = LoadConfig() + model_loader = DiffusersPipelineLoader(load_config) + time_before_load = time.perf_counter() + with DeviceMemoryProfiler() as m: + self.pipeline = model_loader.load_model( + od_config=self.od_config, + load_device=f"cuda:{rank}", ) - - load_config = LoadConfig() - model_loader = DiffusersPipelineLoader(load_config) - time_before_load = time.perf_counter() - with DeviceMemoryProfiler() as m: - self.pipeline = model_loader.load_model( - od_config=self.od_config, - load_device=f"cuda:{rank}", - ) - time_after_load = time.perf_counter() + time_after_load = time.perf_counter() logger.info( "Model loading took %.4f GiB and %.6f seconds", @@ -116,8 +119,8 @@ class GPUWorker: # Refresh cache context if needed if self.cache_backend is not None and self.cache_backend.is_enabled(): self.cache_backend.refresh(self.pipeline, req.num_inference_steps) - - output = self.pipeline.forward(req) + with set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=self.od_config): + output = self.pipeline.forward(req) return output def shutdown(self) -> None: