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

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

@@ -0,0 +1,299 @@
+import torch
+import functools
+from typing import Optional, Tuple
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.transformers.transformer_cogview3plus import (
+    CogView3PlusTransformer2DModel,
+    CogVideoXAttnProcessor2_0,
+)
+from diffusers.models.transformers.transformer_cogview4 import (
+    CogView4Transformer2DModel,
+    CogView4AttnProcessor,
+)
+from diffusers.models.attention_processor import Attention
+from diffusers.models.attention_dispatch import dispatch_attention_fn
+
+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 .cp_plan_cogvideox import __patch_CogVideoXAttnProcessor2_0__call__
+
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+@ContextParallelismPlannerRegister.register("CogView3Plus")
+class CogView3PlusContextParallelismPlanner(ContextParallelismPlanner):
+    def apply(
+        self,
+        transformer: Optional[torch.nn.Module | ModelMixin] = None,
+        **kwargs,
+    ) -> ContextParallelModelPlan:
+
+        # NOTE: Diffusers native CP plan still not supported
+        # for CogView3Plus now.
+        self._cp_planner_preferred_native_diffusers = False
+
+        if (
+            transformer is not None
+            and self._cp_planner_preferred_native_diffusers
+        ):
+            assert isinstance(
+                transformer, CogView3PlusTransformer2DModel
+            ), "Transformer must be an instance of CogView3PlusTransformer2DModel"
+            if hasattr(transformer, "_cp_plan"):
+                if transformer._cp_plan is not None:
+                    return transformer._cp_plan
+
+        # CogView3Plus and CogVideoX share the same attention processor
+        CogVideoXAttnProcessor2_0.__call__ = (
+            __patch_CogVideoXAttnProcessor2_0__call__
+        )
+        # Also need to patch the parallel config and attention backend
+        if not hasattr(CogVideoXAttnProcessor2_0, "_parallel_config"):
+            CogVideoXAttnProcessor2_0._parallel_config = None
+        if not hasattr(CogVideoXAttnProcessor2_0, "_attention_backend"):
+            CogVideoXAttnProcessor2_0._attention_backend = None
+
+        # Otherwise, use the custom CP plan defined here, this maybe
+        # a little different from the native diffusers implementation
+        # for some models.
+        _cp_plan = {
+            # Pattern of transformer_blocks.0, split_output=False:
+            #     un-split input -> split -> to_qkv/...
+            #     -> all2all
+            #     -> attn (local head, full seqlen)
+            #     -> all2all
+            #     -> splited output
+            # Pattern of the rest transformer_blocks, split_output=False:
+            #     splited input (previous splited output) -> to_qkv/...
+            #     -> all2all
+            #     -> attn (local head, full seqlen)
+            #     -> all2all
+            #     -> splited output
+            # The `encoder_hidden_states` will be changed after each block forward,
+            # so we need to split it at the first block, and keep it splited (namely,
+            # automatically split by the all2all op after attn) for the rest blocks.
+            # The `out` tensor of local attn will be splited into `hidden_states` and
+            # `encoder_hidden_states` after each block forward, thus both of them
+            # will be automatically splited by all2all comm op after local attn.
+            "transformer_blocks.0": {
+                "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
+                ),
+            },
+            # transformer forward while using CP, since it is not splited here.
+            # 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
+
+
+@ContextParallelismPlannerRegister.register("CogView4")
+class CogView4ContextParallelismPlanner(ContextParallelismPlanner):
+    def apply(
+        self,
+        transformer: Optional[torch.nn.Module | ModelMixin] = None,
+        **kwargs,
+    ) -> ContextParallelModelPlan:
+
+        # NOTE: Diffusers native CP plan still not supported
+        # for CogView4 now.
+        self._cp_planner_preferred_native_diffusers = False
+
+        if (
+            transformer is not None
+            and self._cp_planner_preferred_native_diffusers
+        ):
+            assert isinstance(
+                transformer, CogView4Transformer2DModel
+            ), "Transformer must be an instance of CogView4Transformer2DModel"
+            if hasattr(transformer, "_cp_plan"):
+                if transformer._cp_plan is not None:
+                    return transformer._cp_plan
+
+        CogView4AttnProcessor.__call__ = __patch_CogView4AttnProcessor__call__
+        # Also need to patch the parallel config and attention backend
+        if not hasattr(CogView4AttnProcessor, "_parallel_config"):
+            CogView4AttnProcessor._parallel_config = None
+        if not hasattr(CogView4AttnProcessor, "_attention_backend"):
+            CogView4AttnProcessor._attention_backend = None
+
+        # Otherwise, use the custom CP plan defined here, this maybe
+        # a little different from the native diffusers implementation
+        # for some models.
+        _cp_plan = {
+            # Pattern of transformer_blocks.0, split_output=False:
+            #     un-split input -> split -> to_qkv/...
+            #     -> all2all
+            #     -> attn (local head, full seqlen)
+            #     -> all2all
+            #     -> splited output
+            # Pattern of the rest transformer_blocks, split_output=False:
+            #     splited input (previous splited output) -> to_qkv/...
+            #     -> all2all
+            #     -> attn (local head, full seqlen)
+            #     -> all2all
+            #     -> splited output
+            # The `encoder_hidden_states` will be changed after each block forward,
+            # so we need to split it at the first block, and keep it splited (namely,
+            # automatically split by the all2all op after attn) for the rest blocks.
+            # The `out` tensor of local attn will be splited into `hidden_states` and
+            # `encoder_hidden_states` after each block forward, thus both of them
+            # will be automatically splited by all2all comm op after local attn.
+            "transformer_blocks.0": {
+                "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
+                ),
+            },
+            # Pattern of the image_rotary_emb, split at every block, because the it
+            # is not automatically splited by all2all comm op and keep un-splited
+            # while the block forward finished:
+            #    un-split input -> split output
+            #    -> after block forward
+            #    -> un-split input
+            #    un-split input -> split output
+            #    ...
+            "transformer_blocks.*": {
+                "image_rotary_emb": [
+                    ContextParallelInput(
+                        split_dim=0, expected_dims=2, split_output=False
+                    ),
+                    ContextParallelInput(
+                        split_dim=0, expected_dims=2, split_output=False
+                    ),
+                ],
+            },
+            # transformer forward while using CP, since it is not splited here.
+            # 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
+
+
+@functools.wraps(CogView4AttnProcessor.__call__)
+def __patch_CogView4AttnProcessor__call__(
+    self: CogView4AttnProcessor,
+    attn: Attention,
+    hidden_states: torch.Tensor,
+    encoder_hidden_states: torch.Tensor,
+    attention_mask: Optional[torch.Tensor] = None,
+    image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    dtype = encoder_hidden_states.dtype
+
+    batch_size, text_seq_length, embed_dim = encoder_hidden_states.shape
+    batch_size, image_seq_length, embed_dim = hidden_states.shape
+    hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+    # 1. QKV projections
+    query = attn.to_q(hidden_states)
+    key = attn.to_k(hidden_states)
+    value = attn.to_v(hidden_states)
+
+    # NOTE(DefTruth): no transpose
+    query = query.unflatten(2, (attn.heads, -1))
+    key = key.unflatten(2, (attn.heads, -1))
+    value = value.unflatten(2, (attn.heads, -1))
+
+    # 2. QK normalization
+    if attn.norm_q is not None:
+        query = attn.norm_q(query).to(dtype=dtype)
+    if attn.norm_k is not None:
+        key = attn.norm_k(key).to(dtype=dtype)
+
+    # 3. Rotational positional embeddings applied to latent stream
+    if image_rotary_emb is not None:
+        from diffusers.models.embeddings import apply_rotary_emb
+
+        query[:, text_seq_length:] = apply_rotary_emb(
+            query[:, text_seq_length:],
+            image_rotary_emb,
+            use_real_unbind_dim=-2,
+            sequence_dim=1,
+        )
+        key[:, text_seq_length:] = apply_rotary_emb(
+            key[:, text_seq_length:],
+            image_rotary_emb,
+            use_real_unbind_dim=-2,
+            sequence_dim=1,
+        )
+
+    # 4. Attention
+    if attention_mask is not None:
+        text_attn_mask = attention_mask
+        assert (
+            text_attn_mask.dim() == 2
+        ), "the shape of text_attn_mask should be (batch_size, text_seq_length)"
+        text_attn_mask = text_attn_mask.float().to(query.device)
+        mix_attn_mask = torch.ones(
+            (batch_size, text_seq_length + image_seq_length),
+            device=query.device,
+        )
+        mix_attn_mask[:, :text_seq_length] = text_attn_mask  # [B, seq_len]
+        # TODO(DefTruth): Permute mix_attn_mask if context parallel is used.
+        # For example, if work size = 2: [E, H] -> [E_0, H_0, E_1, H_1]
+        mix_attn_mask = mix_attn_mask.unsqueeze(2)  # [B, seq_len, 1]
+        attn_mask_matrix = mix_attn_mask @ mix_attn_mask.transpose(
+            1, 2
+        )  # [B, seq_len, seq_len]
+        attention_mask = (
+            (attn_mask_matrix > 0).unsqueeze(1).to(query.dtype)
+        )  # [B, 1, seq_len, seq_len]
+        if (
+            hasattr(self, "_parallel_config")
+            and self._parallel_config is not None
+        ):
+            raise NotImplementedError(
+                "Attention mask with context parallelism for CogView4 "
+                "is not implemented yet."
+            )
+
+    # NOTE(DefTruth): Apply dispatch_attention_fn instead of sdpa directly
+    hidden_states = dispatch_attention_fn(
+        query,
+        key,
+        value,
+        attn_mask=attention_mask,
+        dropout_p=0.0,
+        is_causal=False,
+        backend=getattr(self, "_attention_backend", None),
+        parallel_config=getattr(self, "_parallel_config", None),
+    )
+    hidden_states = hidden_states.flatten(2, 3)
+    hidden_states = hidden_states.type_as(query)
+
+    # 5. Output projection
+    hidden_states = attn.to_out[0](hidden_states)
+    hidden_states = attn.to_out[1](hidden_states)
+
+    encoder_hidden_states, hidden_states = hidden_states.split(
+        [text_seq_length, hidden_states.size(1) - text_seq_length], dim=1
+    )
+    return hidden_states, encoder_hidden_states

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

@@ -0,0 +1,123 @@
+import torch
+from typing import Optional
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.transformers.consisid_transformer_3d import (
+    ConsisIDTransformer3DModel,
+)
+from diffusers.models.transformers.cogvideox_transformer_3d import (
+    CogVideoXAttnProcessor2_0,
+)
+
+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 .cp_plan_cogvideox import __patch_CogVideoXAttnProcessor2_0__call__
+
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+@ContextParallelismPlannerRegister.register("ConsisID")
+class CosisIDContextParallelismPlanner(ContextParallelismPlanner):
+    def apply(
+        self,
+        transformer: Optional[torch.nn.Module | ModelMixin] = None,
+        **kwargs,
+    ) -> ContextParallelModelPlan:
+
+        # NOTE: Diffusers native CP plan still not supported
+        # for ConsisID now.
+        self._cp_planner_preferred_native_diffusers = False
+
+        if (
+            transformer is not None
+            and self._cp_planner_preferred_native_diffusers
+        ):
+            assert isinstance(
+                transformer, ConsisIDTransformer3DModel
+            ), "Transformer must be an instance of ConsisIDTransformer3DModel"
+            if hasattr(transformer, "_cp_plan"):
+                if transformer._cp_plan is not None:
+                    return transformer._cp_plan
+
+        # ConsisID uses the same attention processor as CogVideoX.
+        CogVideoXAttnProcessor2_0.__call__ = (
+            __patch_CogVideoXAttnProcessor2_0__call__
+        )
+        # Also need to patch the parallel config and attention backend
+        if not hasattr(CogVideoXAttnProcessor2_0, "_parallel_config"):
+            CogVideoXAttnProcessor2_0._parallel_config = None
+        if not hasattr(CogVideoXAttnProcessor2_0, "_attention_backend"):
+            CogVideoXAttnProcessor2_0._attention_backend = None
+
+        # Otherwise, use the custom CP plan defined here, this maybe
+        # a little different from the native diffusers implementation
+        # for some models.
+        _cp_plan = {
+            # Pattern of transformer_blocks.0, split_output=False:
+            #     un-split input -> split -> to_qkv/...
+            #     -> all2all
+            #     -> attn (local head, full seqlen)
+            #     -> all2all
+            #     -> splited output
+            # Pattern of the rest transformer_blocks, split_output=False:
+            #     splited input (previous splited output) -> to_qkv/...
+            #     -> all2all
+            #     -> attn (local head, full seqlen)
+            #     -> all2all
+            #     -> splited output
+            # The `encoder_hidden_states` will be changed after each block forward,
+            # so we need to split it at the first block, and keep it splited (namely,
+            # automatically split by the all2all op after attn) for the rest blocks.
+            # The `out` tensor of local attn will be splited into `hidden_states` and
+            # `encoder_hidden_states` after each block forward, thus both of them
+            # will be automatically splited by all2all comm op after local attn.
+            "transformer_blocks.0": {
+                "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
+                ),
+            },
+            # Pattern of the image_rotary_emb, split at every block, because the it
+            # is not automatically splited by all2all comm op and keep un-splited
+            # while the block forward finished:
+            #    un-split input -> split output
+            #    -> after block forward
+            #    -> un-split input
+            #    un-split input -> split output
+            #    ...
+            "transformer_blocks.*": {
+                "image_rotary_emb": [
+                    ContextParallelInput(
+                        split_dim=0, expected_dims=2, split_output=False
+                    ),
+                    ContextParallelInput(
+                        split_dim=0, expected_dims=2, split_output=False
+                    ),
+                ],
+            },
+            # NOTE: We should gather both hidden_states and encoder_hidden_states
+            # at the end of the last block. Because the subsequent op is:
+            # hidden_states = torch.cat([encoder_hidden_states, hidden_states])
+            f"transformer_blocks.{len(transformer.transformer_blocks) - 1}": [
+                ContextParallelOutput(gather_dim=1, expected_dims=3),
+                ContextParallelOutput(gather_dim=1, expected_dims=3),
+            ],
+        }
+        return _cp_plan

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

@@ -0,0 +1,94 @@
+import torch
+from typing import Optional
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.transformers.dit_transformer_2d import (
+    DiTTransformer2DModel,
+)
+from diffusers.models.attention_processor import (
+    Attention,
+    AttnProcessor2_0,
+)  # sdpa
+
+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 .cp_plan_pixart import (
+    __patch_AttnProcessor2_0__call__,
+    __patch_Attention_prepare_attention_mask__,
+)
+
+
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+@ContextParallelismPlannerRegister.register("DiT")
+class DiTContextParallelismPlanner(ContextParallelismPlanner):
+    def apply(
+        self,
+        transformer: Optional[torch.nn.Module | ModelMixin] = None,
+        **kwargs,
+    ) -> ContextParallelModelPlan:
+        assert transformer is not None, "Transformer must be provided."
+        assert isinstance(
+            transformer, DiTTransformer2DModel
+        ), "Transformer must be an instance of DiTTransformer2DModel"
+
+        self._cp_planner_preferred_native_diffusers = False
+
+        if (
+            transformer is not None
+            and self._cp_planner_preferred_native_diffusers
+        ):
+            if hasattr(transformer, "_cp_plan"):
+                if transformer._cp_plan is not None:
+                    return transformer._cp_plan
+
+        # Apply monkey patch to fix attention mask preparation at class level
+        Attention.prepare_attention_mask = (
+            __patch_Attention_prepare_attention_mask__
+        )
+        AttnProcessor2_0.__call__ = __patch_AttnProcessor2_0__call__
+        if not hasattr(AttnProcessor2_0, "_parallel_config"):
+            AttnProcessor2_0._parallel_config = None
+        if not hasattr(AttnProcessor2_0, "_attention_backend"):
+            AttnProcessor2_0._attention_backend = None
+
+        # Otherwise, use the custom CP plan defined here, this maybe
+        # a little different from the native diffusers implementation
+        # for some models.
+
+        _cp_plan = {
+            # Pattern of transformer_blocks.0, split_output=False:
+            #     un-split input -> split -> to_qkv/...
+            #     -> all2all
+            #     -> attn (local head, full seqlen)
+            #     -> all2all
+            #     -> splited output
+            #     (only split hidden_states, not encoder_hidden_states)
+            "transformer_blocks.0": {
+                "hidden_states": ContextParallelInput(
+                    split_dim=1, expected_dims=3, 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_2": ContextParallelOutput(gather_dim=1, expected_dims=3),
+        }
+        return _cp_plan

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

@@ -0,0 +1,88 @@
+import torch
+from typing import Optional
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers import FluxTransformer2DModel
+
+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("Flux")
+class FluxContextParallelismPlanner(ContextParallelismPlanner):
+    def apply(
+        self,
+        transformer: Optional[torch.nn.Module | ModelMixin] = None,
+        **kwargs,
+    ) -> ContextParallelModelPlan:
+        if (
+            transformer is not None
+            and self._cp_planner_preferred_native_diffusers
+        ):
+            assert isinstance(
+                transformer, FluxTransformer2DModel
+            ), "Transformer must be an instance of FluxTransformer2DModel"
+            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 Flux, 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