feat: cp support in qwenimage 2.1

src/diffusers/models/transformers/transformer_qwenimage21.py

@@ -16,6 +16,7 @@
 from typing import Any
 
 import torch
+import torch.distributed as dist
 import torch.nn as nn
 import torch.nn.functional as F
 
@@ -24,6 +25,7 @@
 from ...utils import logging
 from ...utils.peft_utils import apply_lora_scale
 from ...utils.torch_utils import maybe_allow_in_graph
+from .._modeling_parallel import ContextParallelInput, ContextParallelOutput, gather_size_by_comm
 from ..attention import AttentionMixin, AttentionModuleMixin
 from ..attention_dispatch import dispatch_attention_fn
 from ..cache_utils import CacheMixin
@@ -324,13 +326,45 @@ def _qwenimage21_prefix_segments(image_ids: torch.Tensor, prefix_len: int) -> li
     return segments
 
 
+def _qwenimage21_dense_block_causal_mask(
+    segments: list[tuple[int, int, bool]],
+    seq_len: int,
+    key_valid: torch.Tensor | None,
+    batch_size: int,
+    device: torch.device,
+) -> torch.Tensor:
+    attention_mask = torch.ones(seq_len, seq_len, dtype=torch.bool, device=device).tril()
+    for start, end, is_text in segments:
+        if not is_text:
+            attention_mask[start:end, start:end] = True
+    prefix_len = segments[-1][1] if segments else 0
+    attention_mask[prefix_len:] = True
+    attention_mask = attention_mask.view(1, 1, seq_len, seq_len)
+    attention_mask = attention_mask.expand(batch_size, -1, -1, -1)
+    if key_valid is not None:
+        attention_mask = attention_mask & key_valid[:, None, None, :]
+    return attention_mask
+
+
+def _qwenimage21_all_gather_sequence(tensor: torch.Tensor, group: dist.ProcessGroup) -> torch.Tensor:
+    local_sizes = gather_size_by_comm(tensor.shape[1], group)
+    max_local_size = max(local_sizes)
+    if tensor.shape[1] < max_local_size:
+        padding = tensor.new_zeros(tensor.shape[0], max_local_size - tensor.shape[1], *tensor.shape[2:])
+        tensor = torch.cat([tensor, padding], dim=1)
+    gathered = [torch.empty_like(tensor) for _ in local_sizes]
+    dist.all_gather(gathered, tensor, group=group)
+    return torch.cat([value[:, :size] for value, size in zip(gathered, local_sizes)], dim=1)
+
+
 def _qwenimage21_prepare_qkv(
     attn: "QwenImage21Attention",
     hidden_states: torch.Tensor,
     rotary_emb: torch.Tensor | None,
     layer_cache: QwenImage21KVLayerCache | None,
     kv_cache_mode: str | None,
     cache_write_slice: slice | None,
+    parallel_config: Any | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
     """Shared QKV projection, norm, RoPE and KV-cache bookkeeping for both processors."""
     query = attn.to_q(hidden_states)
@@ -350,13 +384,40 @@ def _qwenimage21_prepare_qkv(
 
     if layer_cache is not None:
         if kv_cache_mode == "extract" and cache_write_slice is not None:
-            # `clone()`, not `contiguous()`: at batch size 1 the prefix slice already counts as contiguous
-            # (size-1 dims are ignored), so `contiguous()` returns the same view and the cache would pin the
-            # whole prefill K/V for every step of the denoising loop.
-            layer_cache.store(
-                key[:, cache_write_slice].clone(),
-                value[:, cache_write_slice].clone(),
-            )
+            context_parallel_config = None if parallel_config is None else parallel_config.context_parallel_config
+            if context_parallel_config is None:
+                # `clone()`, not `contiguous()`: at batch size 1 the prefix slice already counts as contiguous
+                # (size-1 dims are ignored), so `contiguous()` returns the same view and the cache would pin the
+                # whole prefill K/V for every step of the denoising loop.
+                layer_cache.store(
+                    key[:, cache_write_slice].clone(),
+                    value[:, cache_write_slice].clone(),
+                )
+            else:
+                group = context_parallel_config._ulysses_mesh.get_group()
+                rank = dist.get_rank(group)
+                world_size = dist.get_world_size(group)
+                local_seq_lens = gather_size_by_comm(key.shape[1], group)
+                local_offset = sum(local_seq_lens[:rank])
+                cache_start = 0 if cache_write_slice.start is None else cache_write_slice.start
+                cache_stop = sum(local_seq_lens) if cache_write_slice.stop is None else cache_write_slice.stop
+                local_cache_start = max(cache_start - local_offset, 0)
+                local_cache_stop = min(cache_stop - local_offset, key.shape[1])
+                local_cache_start = min(local_cache_start, key.shape[1])
+                local_cache_stop = max(local_cache_stop, local_cache_start)
+
+                cached_key = _qwenimage21_all_gather_sequence(key[:, local_cache_start:local_cache_stop], group)
+                cached_value = _qwenimage21_all_gather_sequence(value[:, local_cache_start:local_cache_stop], group)
+                if not context_parallel_config.ulysses_anything and cached_key.shape[1] % world_size != 0:
+                    raise ValueError(
+                        "The cached prefix length must be divisible by the Ulysses degree. Enable "
+                        "`ulysses_anything=True` to cache an uneven prefix."
+                    )
+                split_fn = torch.tensor_split if context_parallel_config.ulysses_anything else torch.chunk
+                layer_cache.store(
+                    split_fn(cached_key, world_size, dim=1)[rank].clone(),
+                    split_fn(cached_value, world_size, dim=1)[rank].clone(),
+                )
         elif kv_cache_mode == "cached":
             cached_k, cached_v = layer_cache.get()
             key = torch.cat([cached_k, key], dim=1)
@@ -401,8 +462,19 @@ def __call__(
         segments: list[tuple[int, int, bool]] | None = None,
         key_valid: torch.Tensor | None = None,
     ) -> torch.Tensor:
+        if self._parallel_config is not None:
+            raise NotImplementedError(
+                "Context parallelism is not implemented for QwenImage21FlexAttnProcessor. "
+                "Use QwenImage21AttnProcessor instead."
+            )
         query, key, value, seq_len_q = _qwenimage21_prepare_qkv(
-            attn, hidden_states, rotary_emb, layer_cache, kv_cache_mode, cache_write_slice
+            attn,
+            hidden_states,
+            rotary_emb,
+            layer_cache,
+            kv_cache_mode,
+            cache_write_slice,
+            self._parallel_config,
         )
 
         seq_len_kv = key.shape[1]
@@ -483,12 +555,34 @@ def __call__(
         segments: list[tuple[int, int, bool]] | None = None,
         key_valid: torch.Tensor | None = None,
     ) -> torch.Tensor:
+        context_parallel_config = (
+            None if self._parallel_config is None else self._parallel_config.context_parallel_config
+        )
+        if context_parallel_config is not None and context_parallel_config.ring_degree > 1:
+            raise NotImplementedError("QwenImage21AttnProcessor currently supports Ulysses context parallelism only.")
+
         query, key, value, seq_len_q = _qwenimage21_prepare_qkv(
-            attn, hidden_states, rotary_emb, layer_cache, kv_cache_mode, cache_write_slice
+            attn,
+            hidden_states,
+            rotary_emb,
+            layer_cache,
+            kv_cache_mode,
+            cache_write_slice,
+            self._parallel_config,
         )
 
         if segments is None:
             # decode: full attention over [cached prefix, target]
+            if context_parallel_config is not None and attention_mask is not None:
+                group = context_parallel_config._ulysses_mesh.get_group()
+                rank = dist.get_rank(group)
+                world_size = dist.get_world_size(group)
+                target_len = sum(gather_size_by_comm(seq_len_q, group))
+                prefix_len = attention_mask.shape[-1] - target_len
+                split_fn = torch.tensor_split if context_parallel_config.ulysses_anything else torch.chunk
+                prefix_masks = split_fn(attention_mask[..., :prefix_len], world_size, dim=-1)
+                target_masks = split_fn(attention_mask[..., prefix_len:], world_size, dim=-1)
+                attention_mask = torch.cat([prefix_masks[rank], target_masks[rank]], dim=-1)
             hidden_states = dispatch_attention_fn(
                 query,
                 key,
@@ -498,6 +592,25 @@ def __call__(
                 backend=self._attention_backend,
                 parallel_config=self._parallel_config,
             )
+        elif context_parallel_config is not None:
+            group = context_parallel_config._ulysses_mesh.get_group()
+            global_seq_len = sum(gather_size_by_comm(seq_len_q, group))
+            attention_mask = _qwenimage21_dense_block_causal_mask(
+                segments,
+                global_seq_len,
+                key_valid,
+                query.shape[0],
+                query.device,
+            )
+            hidden_states = dispatch_attention_fn(
+                query,
+                key,
+                value,
+                attn_mask=attention_mask,
+                dropout_p=0.0,
+                backend=None,
+                parallel_config=self._parallel_config,
+            )
         else:
             # prefill: every segment attends to the keys `[0, end)` (everything before it plus its own block); text
             # segments additionally get a causal triangle over their own keys; padded text keys are dropped.
@@ -758,6 +871,19 @@ class QwenImage21Transformer2DModel(
     _skip_layerwise_casting_patterns = ["pos_embed", "norm"]
     _repeated_blocks = ["QwenImage21TransformerBlock"]
     _skip_keys = ["kv_cache"]
+    _cp_plan = {
+        "transformer_blocks.0": {
+            "hidden_states": ContextParallelInput(split_dim=1, expected_dims=3, split_output=False),
+        },
+        "transformer_blocks.*": {
+            "rotary_emb": ContextParallelInput(split_dim=0, expected_dims=2, split_output=False),
+            "target_token_mask": ContextParallelInput(split_dim=0, expected_dims=1, split_output=False),
+        },
+        "norm_out": {
+            "target_token_mask": ContextParallelInput(split_dim=0, expected_dims=1, split_output=False),
+        },
+        "proj_out": ContextParallelOutput(gather_dim=1, expected_dims=3),
+    }
 
     @register_to_config
     def __init__(

tests/models/transformers/test_models_transformer_qwenimage21.py

@@ -21,10 +21,17 @@
 from diffusers.models.transformers.transformer_qwenimage21 import build_qwenimage21_block_causal_mask
 from diffusers.utils.torch_utils import randn_tensor
 
-from ...testing_utils import enable_full_determinism, torch_device
+from ...testing_utils import (
+    enable_full_determinism,
+    is_context_parallel,
+    require_torch_multi_accelerator,
+    torch_device,
+)
 from ..testing_utils import (
     AttentionTesterMixin,
     BaseModelTesterConfig,
+    ContextParallelAttentionBackendsTesterMixin,
+    ContextParallelTesterMixin,
     MemoryTesterMixin,
     ModelTesterMixin,
     TrainingTesterMixin,
@@ -294,3 +301,42 @@ def test_gradient_checkpointing_is_applied(self):
 
 class TestQwenImage21TransformerAttention(QwenImage21TransformerTesterConfig, AttentionTesterMixin):
     pass
+
+
+@is_context_parallel
+@require_torch_multi_accelerator
+class TestQwenImage21TransformerContextParallel(QwenImage21TransformerTesterConfig, ContextParallelTesterMixin):
+    @pytest.mark.parametrize("cp_type", ["ulysses_degree"], ids=["ulysses"])
+    def test_context_parallel_inference(self, cp_type, batch_size: int = 1):
+        super().test_context_parallel_inference(cp_type, batch_size=batch_size)
+
+    @pytest.mark.parametrize("cp_type", ["ulysses_degree"], ids=["ulysses"])
+    def test_context_parallel_batch_inputs(self, cp_type):
+        super().test_context_parallel_inference(cp_type, batch_size=2)
+
+    @pytest.mark.parametrize("cp_type", ["ulysses_degree"], ids=["ulysses"])
+    def test_context_parallel_backward(self, cp_type, batch_size: int = 1):
+        super().test_context_parallel_backward(cp_type, batch_size=batch_size)
+
+    @pytest.mark.parametrize("cp_type", ["ulysses_degree"], ids=["ulysses"])
+    def test_context_parallel_backward_batch_inputs(self, cp_type):
+        super().test_context_parallel_backward(cp_type, batch_size=2)
+
+    @pytest.mark.parametrize(
+        "cp_type,mesh_shape,mesh_dim_names",
+        [("ulysses_degree", (1, 2, 1), ("ring", "ulysses", "fsdp"))],
+        ids=["ulysses-3d-fsdp"],
+    )
+    def test_context_parallel_custom_mesh(self, cp_type, mesh_shape, mesh_dim_names):
+        super().test_context_parallel_custom_mesh(cp_type, mesh_shape, mesh_dim_names)
+
+
+class TestQwenImage21TransformerContextParallelAttnBackends(
+    QwenImage21TransformerTesterConfig, ContextParallelAttentionBackendsTesterMixin
+):
+    unsupported_attn_backends = ["flash_hub", "flash_varlen_hub", "_flash_3_hub", "_flash_3_varlen_hub"]
+
+    def get_dummy_inputs(self, batch_size: int = 1) -> dict[str, torch.Tensor]:
+        inputs = super().get_dummy_inputs(batch_size=batch_size)
+        inputs["encoder_hidden_states_mask"][:, 1] = 0
+        return inputs