diff --git a/src/diffusers/models/transformers/transformer_qwenimage21.py b/src/diffusers/models/transformers/transformer_qwenimage21.py index b6eabfd584bb..f5887bb0c6b6 100644 --- a/src/diffusers/models/transformers/transformer_qwenimage21.py +++ b/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,6 +326,37 @@ 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, @@ -331,6 +364,7 @@ def _qwenimage21_prepare_qkv( 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__( diff --git a/tests/models/transformers/test_models_transformer_qwenimage21.py b/tests/models/transformers/test_models_transformer_qwenimage21.py index cb23fc4ad1d2..b3de7b19bf31 100644 --- a/tests/models/transformers/test_models_transformer_qwenimage21.py +++ b/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