cache-dit 1.0.3__py3-none-any.whl → 1.0.14__py3-none-any.whl

cache_dit/parallelism/backends/native_diffusers/context_parallelism/__init__.py

@@ -0,0 +1,164 @@
+import torch
+from typing import Optional
+
+from diffusers.models.modeling_utils import ModelMixin
+from cache_dit.parallelism.parallel_backend import ParallelismBackend
+from cache_dit.parallelism.parallel_config import ParallelismConfig
+from cache_dit.logger import init_logger
+from ..utils import (
+    native_diffusers_parallelism_available,
+    ContextParallelConfig,
+)
+from .attention import maybe_resigter_native_attention_backend
+from .cp_planners import *
+
+try:
+    maybe_resigter_native_attention_backend()
+except ImportError as e:
+    raise ImportError(e)
+
+logger = init_logger(__name__)
+
+
+def maybe_enable_context_parallelism(
+    transformer: torch.nn.Module,
+    parallelism_config: Optional[ParallelismConfig],
+) -> torch.nn.Module:
+    assert isinstance(transformer, ModelMixin), (
+        "transformer must be an instance of diffusers' ModelMixin, "
+        f"but got {type(transformer)}"
+    )
+    if parallelism_config is None:
+        return transformer
+
+    assert isinstance(parallelism_config, ParallelismConfig), (
+        "parallelism_config must be an instance of ParallelismConfig"
+        f" but got {type(parallelism_config)}"
+    )
+
+    if (
+        parallelism_config.backend == ParallelismBackend.NATIVE_DIFFUSER
+        and native_diffusers_parallelism_available()
+    ):
+        cp_config = None
+        if (
+            parallelism_config.ulysses_size is not None
+            or parallelism_config.ring_size is not None
+        ):
+            cp_config = ContextParallelConfig(
+                ulysses_degree=parallelism_config.ulysses_size,
+                ring_degree=parallelism_config.ring_size,
+            )
+        if cp_config is not None:
+            attention_backend = parallelism_config.parallel_kwargs.get(
+                "attention_backend", None
+            )
+            if hasattr(transformer, "enable_parallelism"):
+                if hasattr(transformer, "set_attention_backend"):
+                    # native, _native_cudnn, flash, etc.
+                    if attention_backend is None:
+                        # Now only _native_cudnn is supported for parallelism
+                        # issue: https://github.com/huggingface/diffusers/pull/12443
+                        transformer.set_attention_backend("_native_cudnn")
+                        logger.warning(
+                            "attention_backend is None, set default attention backend "
+                            "to _native_cudnn for parallelism because of the issue: "
+                            "https://github.com/huggingface/diffusers/pull/12443"
+                        )
+                    else:
+                        transformer.set_attention_backend(attention_backend)
+                        logger.info(
+                            "Found attention_backend from config, set attention "
+                            f"backend to: {attention_backend}"
+                        )
+                # Prefer custom cp_plan if provided
+                cp_plan = parallelism_config.parallel_kwargs.get(
+                    "cp_plan", None
+                )
+                if cp_plan is not None:
+                    logger.info(
+                        f"Using custom context parallelism plan: {cp_plan}"
+                    )
+                else:
+                    # Try get context parallelism plan from register if not provided
+                    extra_parallel_kwargs = {}
+                    if parallelism_config.parallel_kwargs is not None:
+                        extra_parallel_kwargs = (
+                            parallelism_config.parallel_kwargs
+                        )
+                    cp_plan = ContextParallelismPlannerRegister.get_planner(
+                        transformer
+                    )().apply(transformer=transformer, **extra_parallel_kwargs)
+
+                transformer.enable_parallelism(
+                    config=cp_config, cp_plan=cp_plan
+                )
+                _maybe_patch_native_parallel_config(transformer)
+            else:
+                raise ValueError(
+                    f"{transformer.__class__.__name__} does not support context parallelism."
+                )
+
+    return transformer
+
+
+def _maybe_patch_native_parallel_config(
+    transformer: torch.nn.Module,
+) -> torch.nn.Module:
+
+    cls_name = transformer.__class__.__name__
+    if not cls_name.startswith("Nunchaku"):
+        return transformer
+
+    from diffusers import FluxTransformer2DModel, QwenImageTransformer2DModel
+
+    try:
+        from nunchaku.models.transformers.transformer_flux_v2 import (
+            NunchakuFluxTransformer2DModelV2,
+            NunchakuFluxAttention,
+            NunchakuFluxFA2Processor,
+        )
+        from nunchaku.models.transformers.transformer_qwenimage import (
+            NunchakuQwenAttention,
+            NunchakuQwenImageNaiveFA2Processor,
+            NunchakuQwenImageTransformer2DModel,
+        )
+    except ImportError:
+        raise ImportError(
+            "NunchakuFluxTransformer2DModelV2 or NunchakuQwenImageTransformer2DModel "
+            "requires the 'nunchaku' package. Please install nunchaku before using "
+            "the context parallelism for nunchaku 4-bits models."
+        )
+    assert isinstance(
+        transformer,
+        (
+            NunchakuFluxTransformer2DModelV2,
+            FluxTransformer2DModel,
+        ),
+    ) or isinstance(
+        transformer,
+        (
+            NunchakuQwenImageTransformer2DModel,
+            QwenImageTransformer2DModel,
+        ),
+    ), (
+        "transformer must be an instance of NunchakuFluxTransformer2DModelV2 "
+        f"or NunchakuQwenImageTransformer2DModel, but got {type(transformer)}"
+    )
+    config = transformer._parallel_config
+
+    attention_classes = (
+        NunchakuFluxAttention,
+        NunchakuFluxFA2Processor,
+        NunchakuQwenAttention,
+        NunchakuQwenImageNaiveFA2Processor,
+    )
+    for module in transformer.modules():
+        if not isinstance(module, attention_classes):
+            continue
+        processor = getattr(module, "processor", None)
+        if processor is None or not hasattr(processor, "_parallel_config"):
+            continue
+        processor._parallel_config = config
+
+    return transformer

cache_dit/parallelism/backends/native_diffusers/context_parallelism/attention/__init__.py

@@ -0,0 +1,4 @@
+def maybe_resigter_native_attention_backend():
+    """Maybe re-register native attention backend to enable context parallelism."""
+    # Import custom attention backend ensuring registration
+    from ._attention_dispatch import _native_attention

cache_dit/parallelism/backends/native_diffusers/context_parallelism/attention/_attention_dispatch.py

@@ -0,0 +1,304 @@
+import os
+import torch
+from typing import Optional
+
+try:
+    from diffusers.models.attention_dispatch import (
+        _AttentionBackendRegistry,
+        AttentionBackendName,
+        _check_device,
+        _check_shape,
+        TemplatedRingAttention,
+        TemplatedUlyssesAttention,
+    )
+    from diffusers.models._modeling_parallel import ParallelConfig
+except ImportError:
+    raise ImportError(
+        "Context parallelism requires the 'diffusers>=0.36.dev0'."
+        "Please install latest version of diffusers from source: \n"
+        "pip3 install git+https://github.com/huggingface/diffusers.git"
+    )
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+__all__ = [
+    "_native_attention",
+]
+
+# Enable custom native attention backend with context parallelism
+# by default. Users can set the environment variable to 0 to disable
+# this behavior. Default to enabled for better compatibility.
+_CACHE_DIT_ENABLE_CUSTOM_CP_NATIVE_ATTN_DISPATCH = bool(
+    int(os.getenv("CACHE_DIT_ENABLE_CUSTOM_CP_NATIVE_ATTN_DISPATCH", "1"))
+)
+
+
+def _is_native_attn_supported_context_parallel() -> bool:
+    try:
+        return (
+            AttentionBackendName.NATIVE
+            in _AttentionBackendRegistry._supports_context_parallel
+            and _AttentionBackendRegistry._supports_context_parallel[
+                AttentionBackendName.NATIVE
+            ]
+        )
+    except Exception:
+        assert isinstance(
+            _AttentionBackendRegistry._supports_context_parallel, set
+        )
+        return (
+            AttentionBackendName.NATIVE.value
+            in _AttentionBackendRegistry._supports_context_parallel
+        )
+
+
+if _CACHE_DIT_ENABLE_CUSTOM_CP_NATIVE_ATTN_DISPATCH:
+    logger.warning(
+        "Re-registering NATIVE attention backend to enable context parallelism. "
+        "This is a temporary workaround and should be removed after the native "
+        "attention backend supports context parallelism natively. Please check: "
+        "https://github.com/huggingface/diffusers/pull/12563 for more details. "
+        "Or, you can disable this behavior by setting the environment variable "
+        "`CACHE_DIT_ENABLE_CUSTOM_CP_NATIVE_ATTN_DISPATCH=0`."
+    )
+    _AttentionBackendRegistry._backends.pop(AttentionBackendName.NATIVE)
+    _AttentionBackendRegistry._constraints.pop(AttentionBackendName.NATIVE)
+    _AttentionBackendRegistry._supported_arg_names.pop(
+        AttentionBackendName.NATIVE
+    )
+    if _is_native_attn_supported_context_parallel():
+        if isinstance(
+            _AttentionBackendRegistry._supports_context_parallel, dict
+        ):
+            _AttentionBackendRegistry._supports_context_parallel.pop(
+                AttentionBackendName.NATIVE
+            )
+        else:
+            _AttentionBackendRegistry._supports_context_parallel.remove(
+                AttentionBackendName.NATIVE.value
+            )
+
+    # Re-define templated context parallel attention to support attn mask
+    def _templated_context_parallel_attention_v2(
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,
+        dropout_p: float = 0.0,
+        is_causal: bool = False,
+        scale: Optional[float] = None,
+        enable_gqa: bool = False,
+        return_lse: bool = False,
+        *,
+        forward_op,
+        backward_op,
+        _parallel_config: Optional["ParallelConfig"] = None,
+    ):
+        if attn_mask is not None:
+            # NOTE(DefTruth): Check if forward_op is native attention forward op
+            forward_op_name = forward_op.__name__
+            if not forward_op_name == "_native_attention_forward_op":
+                raise ValueError(
+                    "Templated context parallel attention with attn_mask "
+                    "is only supported for native attention backend, "
+                    f"but got forward_op: {forward_op_name}."
+                )
+        if is_causal:
+            raise ValueError(
+                "Causal attention is not yet supported for templated attention."
+            )
+        if enable_gqa:
+            raise ValueError(
+                "GQA is not yet supported for templated attention."
+            )
+
+        # TODO: add support for unified attention with ring/ulysses degree both being > 1
+        if _parallel_config.context_parallel_config.ring_degree > 1:
+            return TemplatedRingAttention.apply(
+                query,
+                key,
+                value,
+                attn_mask,
+                dropout_p,
+                is_causal,
+                scale,
+                enable_gqa,
+                return_lse,
+                forward_op,
+                backward_op,
+                _parallel_config,
+            )
+        elif _parallel_config.context_parallel_config.ulysses_degree > 1:
+            return TemplatedUlyssesAttention.apply(
+                query,
+                key,
+                value,
+                attn_mask,
+                dropout_p,
+                is_causal,
+                scale,
+                enable_gqa,
+                return_lse,
+                forward_op,
+                backward_op,
+                _parallel_config,
+            )
+        else:
+            raise ValueError(
+                "Reaching this branch of code is unexpected. Please report a bug."
+            )
+
+    # NOTE:Remove NATIVE attention backend constraints and re-register it.
+    # Here is a temporary workaround to enable context parallelism with
+    # native attention backend. We should remove this workaround after
+    # the native attention backend supports context parallelism natively.
+    # Adapted from: https://github.com/huggingface/diffusers/pull/12563
+
+    def _native_attention_forward_op(
+        ctx: torch.autograd.function.FunctionCtx,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,
+        dropout_p: float = 0.0,
+        is_causal: bool = False,
+        scale: Optional[float] = None,
+        enable_gqa: bool = False,
+        return_lse: bool = False,
+        _save_ctx: bool = True,
+        _parallel_config: Optional["ParallelConfig"] = None,
+    ):
+        # Native attention does not return_lse
+        if return_lse:
+            raise ValueError(
+                "Native attention does not support return_lse=True"
+            )
+
+        # used for backward pass
+        if _save_ctx:
+            ctx.save_for_backward(query, key, value)
+            ctx.attn_mask = attn_mask
+            ctx.dropout_p = dropout_p
+            ctx.is_causal = is_causal
+            ctx.scale = scale
+            ctx.enable_gqa = enable_gqa
+
+        query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value))
+        out = torch.nn.functional.scaled_dot_product_attention(
+            query=query,
+            key=key,
+            value=value,
+            attn_mask=attn_mask,
+            dropout_p=dropout_p,
+            is_causal=is_causal,
+            scale=scale,
+            enable_gqa=enable_gqa,
+        )
+        out = out.permute(0, 2, 1, 3)
+
+        return out
+
+    def _native_attention_backward_op(
+        ctx: torch.autograd.function.FunctionCtx,
+        grad_out: torch.Tensor,
+        *args,
+        **kwargs,
+    ):
+        query, key, value = ctx.saved_tensors
+
+        query.requires_grad_(True)
+        key.requires_grad_(True)
+        value.requires_grad_(True)
+
+        query_t, key_t, value_t = (
+            x.permute(0, 2, 1, 3) for x in (query, key, value)
+        )
+        out = torch.nn.functional.scaled_dot_product_attention(
+            query=query_t,
+            key=key_t,
+            value=value_t,
+            attn_mask=ctx.attn_mask,
+            dropout_p=ctx.dropout_p,
+            is_causal=ctx.is_causal,
+            scale=ctx.scale,
+            enable_gqa=ctx.enable_gqa,
+        )
+        out = out.permute(0, 2, 1, 3)
+
+        grad_out_t = grad_out.permute(0, 2, 1, 3)
+        grad_query_t, grad_key_t, grad_value_t = torch.autograd.grad(
+            outputs=out,
+            inputs=[query_t, key_t, value_t],
+            grad_outputs=grad_out_t,
+            retain_graph=False,
+        )
+
+        grad_query = grad_query_t.permute(0, 2, 1, 3)
+        grad_key = grad_key_t.permute(0, 2, 1, 3)
+        grad_value = grad_value_t.permute(0, 2, 1, 3)
+
+        return grad_query, grad_key, grad_value
+
+    @_AttentionBackendRegistry.register(
+        AttentionBackendName.NATIVE,
+        constraints=[_check_device, _check_shape],
+        supports_context_parallel=True,
+    )
+    def _native_attention(
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,
+        dropout_p: float = 0.0,
+        is_causal: bool = False,
+        scale: Optional[float] = None,
+        enable_gqa: bool = False,
+        return_lse: bool = False,
+        _parallel_config: Optional["ParallelConfig"] = None,
+    ) -> torch.Tensor:
+        if return_lse:
+            raise ValueError(
+                "Native attention backend does not support setting `return_lse=True`."
+            )
+        if _parallel_config is None:
+            query, key, value = (
+                x.permute(0, 2, 1, 3) for x in (query, key, value)
+            )
+            out = torch.nn.functional.scaled_dot_product_attention(
+                query=query,
+                key=key,
+                value=value,
+                attn_mask=attn_mask,
+                dropout_p=dropout_p,
+                is_causal=is_causal,
+                scale=scale,
+                enable_gqa=enable_gqa,
+            )
+            out = out.permute(0, 2, 1, 3)
+        else:
+            out = _templated_context_parallel_attention_v2(
+                query,
+                key,
+                value,
+                attn_mask,
+                dropout_p,
+                is_causal,
+                scale,
+                enable_gqa,
+                return_lse,
+                forward_op=_native_attention_forward_op,
+                backward_op=_native_attention_backward_op,
+                _parallel_config=_parallel_config,
+            )
+        return out
+
+else:
+    from diffusers.models.attention_dispatch import (
+        _native_attention,
+    )  # noqa: F401
+
+    logger.info(
+        "Native attention backend already supports context parallelism."
+    )

cache_dit/parallelism/backends/native_diffusers/context_parallelism/cp_plan_chroma.py

@@ -0,0 +1,95 @@
+import torch
+from typing import Optional
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.transformers.transformer_chroma import (
+    ChromaTransformer2DModel,
+)
+
+try:
+    from diffusers.models._modeling_parallel import (
+        ContextParallelInput,
+        ContextParallelOutput,
+        ContextParallelModelPlan,
+    )
+except ImportError:
+    raise ImportError(
+        "Context parallelism requires the 'diffusers>=0.36.dev0'."
+        "Please install latest version of diffusers from source: \n"
+        "pip3 install git+https://github.com/huggingface/diffusers.git"
+    )
+from .cp_plan_registers import (
+    ContextParallelismPlanner,
+    ContextParallelismPlannerRegister,
+)
+
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+@ContextParallelismPlannerRegister.register("Chroma")
+class ChromaContextParallelismPlanner(ContextParallelismPlanner):
+    def apply(
+        self,
+        transformer: Optional[torch.nn.Module | ModelMixin] = None,
+        **kwargs,
+    ) -> ContextParallelModelPlan:
+
+        # NOTE: Diffusers native CP plan still not supported
+        # for Chroma now.
+        self._cp_planner_preferred_native_diffusers = False
+
+        if (
+            transformer is not None
+            and self._cp_planner_preferred_native_diffusers
+        ):
+            assert isinstance(
+                transformer, ChromaTransformer2DModel
+            ), "Transformer must be an instance of ChromaTransformer2DModel"
+            if hasattr(transformer, "_cp_plan"):
+                if transformer._cp_plan is not None:
+                    return transformer._cp_plan
+
+        # Otherwise, use the custom CP plan defined here, this maybe
+        # a little different from the native diffusers implementation
+        # for some models.
+        _cp_plan = {
+            # Here is a Transformer level CP plan for Chroma, which will
+            # only apply the only 1 split hook (pre_forward) on the forward
+            # of Transformer, and gather the output after Transformer forward.
+            # Pattern of transformer forward, split_output=False:
+            #     un-split input -> splited input (inside transformer)
+            # Pattern of the transformer_blocks, single_transformer_blocks:
+            #     splited input (previous splited output) -> to_qkv/...
+            #     -> all2all
+            #     -> attn (local head, full seqlen)
+            #     -> all2all
+            #     -> splited output
+            # The `hidden_states` and `encoder_hidden_states` will still keep
+            # itself splited after block forward (namely, automatic split by
+            # the all2all comm op after attn) for the all blocks.
+            # img_ids and txt_ids will only be splited once at the very beginning,
+            # and keep splited through the whole transformer forward. The all2all
+            # comm op only happens on the `out` tensor after local attn not on
+            # img_ids and txt_ids.
+            "": {
+                "hidden_states": ContextParallelInput(
+                    split_dim=1, expected_dims=3, split_output=False
+                ),
+                "encoder_hidden_states": ContextParallelInput(
+                    split_dim=1, expected_dims=3, split_output=False
+                ),
+                "img_ids": ContextParallelInput(
+                    split_dim=0, expected_dims=2, split_output=False
+                ),
+                "txt_ids": ContextParallelInput(
+                    split_dim=0, expected_dims=2, split_output=False
+                ),
+            },
+            # Then, the final proj_out will gather the splited output.
+            #     splited input (previous splited output)
+            #     -> all gather
+            #     -> un-split output
+            "proj_out": ContextParallelOutput(gather_dim=1, expected_dims=3),
+        }
+        return _cp_plan