[Diffusion][Bug-Fix]: Fix Ulysses SP Accurary Error (#377)

Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com>

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: