cache-dit 1.0.9__py3-none-any.whl → 1.0.10__py3-none-any.whl

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

@@ -0,0 +1,254 @@
+# Docstring references: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/_modeling_parallel.py#L185
+# A dictionary where keys denote the input to be split across context parallel region, and the
+# value denotes the sharding configuration.
+# If the key is a string, it denotes the name of the parameter in the forward function.
+# If the key is an integer, split_output must be set to True, and it denotes the index of the output
+# to be split across context parallel region.
+# ContextParallelInputType = Dict[
+#     Union[str, int], Union[ContextParallelInput, List[ContextParallelInput], Tuple[ContextParallelInput, ...]]
+# ]
+
+# A dictionary where keys denote the output to be gathered across context parallel region, and the
+# value denotes the gathering configuration.
+# ContextParallelOutputType = Union[
+#     ContextParallelOutput, List[ContextParallelOutput], Tuple[ContextParallelOutput, ...]
+# ]
+
+# A dictionary where keys denote the module id, and the value denotes how the inputs/outputs of
+# the module should be split/gathered across context parallel region.
+# ContextParallelModelPlan = Dict[str, Union[ContextParallelInputType, ContextParallelOutputType]]
+
+# Example of a ContextParallelModelPlan (QwenImageTransformer2DModel):
+#
+# Each model should define a _cp_plan attribute that contains information on how to shard/gather
+# tensors at different stages of the forward:
+#
+# ```python
+# _cp_plan = {
+#     "": {
+#         "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),
+#         "encoder_hidden_states_mask": ContextParallelInput(split_dim=1, expected_dims=2, split_output=False),
+#     },
+#     "pos_embed": {
+#         0: ContextParallelInput(split_dim=0, expected_dims=2, split_output=True),
+#         1: ContextParallelInput(split_dim=0, expected_dims=2, split_output=True),
+#     },
+#     "proj_out": ContextParallelOutput(gather_dim=1, expected_dims=3),
+# }
+# ```
+#
+# The dictionary is a set of module names mapped to their respective CP plan. The inputs/outputs of layers will be
+# split/gathered according to this at the respective module level. Here, the following happens:
+# - "":
+#     we specify that we want to split the various inputs across the sequence dim in the pre-forward hook (i.e. before
+#     the actual forward logic of the QwenImageTransformer2DModel is run, we will splitthe inputs)
+# - "pos_embed":
+#     we specify that we want to split the outputs of the RoPE layer. Since there are two outputs (imag & text freqs),
+#     we can individually specify how they should be split
+# - "proj_out":
+#     before returning to the user, we gather the entire sequence on each rank in the post-forward hook (after the linear
+#     layer forward has run).
+#
+# ContextParallelInput:
+#     specifies how to split the input tensor in the pre-forward or post-forward hook of the layer it is attached to
+#
+# ContextParallelOutput:
+#     specifies how to gather the input tensor in the post-forward hook in the layer it is attached to
+
+import torch
+from typing import Optional
+from diffusers.models.modeling_utils import ModelMixin
+
+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__)
+
+
+__all__ = [
+    "ContextParallelismPlanner",
+    "ContextParallelismPlannerRegister",
+    "FluxContextParallelismPlanner",
+    "QwenImageContextParallelismPlanner",
+    "WanContextParallelismPlanner",
+    "LTXVideoContextParallelismPlanner",
+]
+
+
+# Register context parallelism planner for models
+@ContextParallelismPlannerRegister.register("Flux")
+class FluxContextParallelismPlanner(ContextParallelismPlanner):
+    def apply(
+        self,
+        transformer: Optional[torch.nn.Module | ModelMixin] = None,
+        **kwargs,
+    ) -> ContextParallelModelPlan:
+        if transformer is not None:
+            from diffusers import FluxTransformer2DModel
+
+            assert isinstance(
+                transformer, FluxTransformer2DModel
+            ), "Transformer must be an instance of FluxTransformer2DModel"
+            if hasattr(transformer, "_cp_plan"):
+                return transformer._cp_plan
+
+        _cp_plan = {
+            "": {
+                "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
+                ),
+            },
+            "proj_out": ContextParallelOutput(gather_dim=1, expected_dims=3),
+        }
+        return _cp_plan
+
+
+@ContextParallelismPlannerRegister.register("QwenImage")
+class QwenImageContextParallelismPlanner(ContextParallelismPlanner):
+    def apply(
+        self,
+        transformer: Optional[torch.nn.Module | ModelMixin] = None,
+        **kwargs,
+    ) -> ContextParallelModelPlan:
+        if transformer is not None:
+            from diffusers import QwenImageTransformer2DModel
+
+            assert isinstance(
+                transformer, QwenImageTransformer2DModel
+            ), "Transformer must be an instance of QwenImageTransformer2DModel"
+            if hasattr(transformer, "_cp_plan"):
+                return transformer._cp_plan
+
+        _cp_plan = _cp_plan = {
+            "": {
+                "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
+                ),
+                "encoder_hidden_states_mask": ContextParallelInput(
+                    split_dim=1, expected_dims=2, split_output=False
+                ),
+            },
+            "pos_embed": {
+                0: ContextParallelInput(
+                    split_dim=0, expected_dims=2, split_output=True
+                ),
+                1: ContextParallelInput(
+                    split_dim=0, expected_dims=2, split_output=True
+                ),
+            },
+            "proj_out": ContextParallelOutput(gather_dim=1, expected_dims=3),
+        }
+        return _cp_plan
+
+
+# TODO: Add WanVACETransformer3DModel context parallelism planner.
+# NOTE: We choice to use full name to avoid name conflict between
+# WanTransformer3DModel and WanVACETransformer3DModel.
+@ContextParallelismPlannerRegister.register("WanTransformer3D")
+class WanContextParallelismPlanner(ContextParallelismPlanner):
+    def apply(
+        self,
+        transformer: Optional[torch.nn.Module | ModelMixin] = None,
+        **kwargs,
+    ) -> ContextParallelModelPlan:
+        if transformer is not None:
+            from diffusers import WanTransformer3DModel
+
+            assert isinstance(
+                transformer, WanTransformer3DModel
+            ), "Transformer must be an instance of WanTransformer3DModel"
+            if hasattr(transformer, "_cp_plan"):
+                return transformer._cp_plan
+
+        _cp_plan = {
+            "rope": {
+                0: ContextParallelInput(
+                    split_dim=1, expected_dims=4, split_output=True
+                ),
+                1: ContextParallelInput(
+                    split_dim=1, expected_dims=4, split_output=True
+                ),
+            },
+            "blocks.0": {
+                "hidden_states": ContextParallelInput(
+                    split_dim=1, expected_dims=3, split_output=False
+                ),
+            },
+            "blocks.*": {
+                "encoder_hidden_states": ContextParallelInput(
+                    split_dim=1, expected_dims=3, split_output=False
+                ),
+            },
+            "proj_out": ContextParallelOutput(gather_dim=1, expected_dims=3),
+        }
+        return _cp_plan
+
+
+@ContextParallelismPlannerRegister.register("LTXVideo")
+class LTXVideoContextParallelismPlanner(ContextParallelismPlanner):
+    def apply(
+        self,
+        transformer: Optional[torch.nn.Module | ModelMixin] = None,
+        **kwargs,
+    ) -> ContextParallelModelPlan:
+        if transformer is not None:
+            from diffusers import LTXVideoTransformer3DModel
+
+            assert isinstance(
+                transformer, LTXVideoTransformer3DModel
+            ), "Transformer must be an instance of LTXVideoTransformer3DModel"
+            if hasattr(transformer, "_cp_plan"):
+                return transformer._cp_plan
+
+        _cp_plan = {
+            "": {
+                "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
+                ),
+                "encoder_attention_mask": ContextParallelInput(
+                    split_dim=1, expected_dims=2, split_output=False
+                ),
+            },
+            "rope": {
+                0: ContextParallelInput(
+                    split_dim=1, expected_dims=3, split_output=True
+                ),
+                1: ContextParallelInput(
+                    split_dim=1, expected_dims=3, split_output=True
+                ),
+            },
+            "proj_out": ContextParallelOutput(gather_dim=1, expected_dims=3),
+        }
+        return _cp_plan

cache_dit/parallelism/backends/native_diffusers/parallel_difffusers.py

@@ -6,22 +6,10 @@ from cache_dit.logger import init_logger
 logger = init_logger(__name__)
 
 
-try:
-    from diffusers import ContextParallelConfig
-
-    def native_diffusers_parallelism_available() -> bool:
-        return True
-
-except ImportError:
-    ContextParallelConfig = None
-
-    def native_diffusers_parallelism_available() -> bool:
-        return False
-
-
 from diffusers.models.modeling_utils import ModelMixin
 from cache_dit.parallelism.parallel_backend import ParallelismBackend
 from cache_dit.parallelism.parallel_config import ParallelismConfig
+from .context_parallelism import maybe_enable_context_parallelism
 
 
 def maybe_enable_parallelism(
@@ -40,54 +28,22 @@ def maybe_enable_parallelism(
         f" but got {type(parallelism_config)}"
     )
 
+    assert parallelism_config.backend == ParallelismBackend.NATIVE_DIFFUSER, (
+        f"parallelism backend must be {ParallelismBackend.NATIVE_DIFFUSER}, "
+        f"but got {parallelism_config.backend}"
+    )
+
     if (
-        parallelism_config.backend == ParallelismBackend.NATIVE_DIFFUSER
-        and native_diffusers_parallelism_available()
+        parallelism_config.ulysses_size is not None
+        or parallelism_config.ring_size is not None
     ):
-        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_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}"
-                        )
-                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}"
-                    )
-                transformer.enable_parallelism(
-                    config=cp_config, cp_plan=cp_plan
-                )
-            else:
-                raise ValueError(
-                    f"{transformer.__class__.__name__} does not support context parallelism."
-                )
-
+        transformer = maybe_enable_context_parallelism(
+            transformer,
+            parallelism_config,
+        )
+    else:
+        raise ValueError(
+            "NATIVE_DIFFUSER backend only support context parallelism now. "
+            "Please set ulysses_size or ring_size in parallelism_config."
+        )
     return transformer

cache_dit/parallelism/backends/native_diffusers/utils.py

@@ -0,0 +1,11 @@
+try:
+    from diffusers import ContextParallelConfig
+
+    def native_diffusers_parallelism_available() -> bool:
+        return True
+
+except ImportError:
+    ContextParallelConfig = None
+
+    def native_diffusers_parallelism_available() -> bool:
+        return False

cache_dit/parallelism/backends/native_pytorch/__init__.py

@@ -0,0 +1,3 @@
+from cache_dit.parallelism.backends.native_pytorch.parallel_torch import (
+    maybe_enable_parallelism,
+)

cache_dit/parallelism/backends/native_pytorch/parallel_torch.py

@@ -0,0 +1,62 @@
+from typing import Optional
+
+import torch
+
+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
+
+logger = init_logger(__name__)
+
+
+def maybe_enable_parallelism(
+    transformer: torch.nn.Module | ModelMixin,
+    parallelism_config: Optional[ParallelismConfig],
+) -> torch.nn.Module:
+    assert isinstance(transformer, torch.nn.Module), (
+        "transformer must be an instance of torch.nn.Module, "
+        f"but got {type(transformer)}"
+    )
+    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 parallelism_config.backend == ParallelismBackend.NATIVE_PYTORCH, (
+        "parallelism_config.backend must be ParallelismBackend.NATIVE_PYTORCH "
+        f"but got {parallelism_config.backend}"
+    )
+
+    assert isinstance(parallelism_config, ParallelismConfig), (
+        "parallelism_config must be an instance of ParallelismConfig"
+        f" but got {type(parallelism_config)}"
+    )
+    assert (
+        parallelism_config.ulysses_size is None
+        and parallelism_config.ring_size is None
+    ), (
+        "Ulysses/Ring parallelism is not supported in Native_PyTorch backend. "
+        "Please set it to None in parallelism_config."
+    )
+
+    if (
+        parallelism_config.tp_size is not None
+        and parallelism_config.tp_size > 1
+    ):
+        from .tensor_parallelism import maybe_enable_tensor_parallelism
+
+        transformer = maybe_enable_tensor_parallelism(
+            transformer=transformer,
+            parallelism_config=parallelism_config,
+        )
+    else:
+        raise ValueError(
+            "NATIVE_PYTORCH only supported tensor parallelism now. "
+            "Please set tp_size > 1 for tensor parallelism."
+        )
+    return transformer

cache_dit/parallelism/backends/native_pytorch/tensor_parallelism/__init__.py

@@ -0,0 +1,48 @@
+try:
+    import einops
+except ImportError:
+    raise ImportError(
+        "Metrics functionality requires the 'parallelism' extra dependencies. "
+        "Install with:\npip install cache-dit[parallelism]"
+    )
+
+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 .tp_planners import *
+
+logger = init_logger(__name__)
+
+
+def maybe_enable_tensor_parallelism(
+    transformer: torch.nn.Module | ModelMixin,
+    parallelism_config: Optional[ParallelismConfig],
+) -> torch.nn.Module:
+    assert isinstance(transformer, torch.nn.Module), (
+        "transformer must be an instance of torch.nn.Module, "
+        f"but got {type(transformer)}"
+    )
+    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 parallelism_config.backend == ParallelismBackend.NATIVE_PYTORCH, (
+        "parallelism_config.backend must be ParallelismBackend.NATIVE_PYTORCH "
+        f"but got {parallelism_config.backend}"
+    )
+
+    extra_parallel_kwargs = {}
+    if parallelism_config.parallel_kwargs is not None:
+        extra_parallel_kwargs = parallelism_config.parallel_kwargs
+
+    return TensorParallelismPlannerRegister.get_planner(transformer)().apply(
+        transformer=transformer,
+        parallelism_config=parallelism_config,
+        **extra_parallel_kwargs,
+    )

cache_dit/parallelism/backends/native_pytorch/tensor_parallelism/tp_plan_flux.py

@@ -0,0 +1,159 @@
+import torch
+from diffusers.models.transformers.transformer_flux import (
+    FluxSingleTransformerBlock,
+)
+from einops import rearrange
+from torch import nn
+from torch.distributed import DeviceMesh, init_device_mesh
+from torch.distributed._tensor import Replicate
+from torch.distributed.tensor.parallel import (
+    ColwiseParallel,
+    RowwiseParallel,
+    parallelize_module,
+)
+from cache_dit.parallelism.parallel_config import ParallelismConfig
+from .tp_plan_registers import (
+    TensorParallelismPlanner,
+    TensorParallelismPlannerRegister,
+)
+
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+@TensorParallelismPlannerRegister.register("Flux")
+class FluxTensorParallelismPlanner(TensorParallelismPlanner):
+    def apply(
+        self,
+        transformer: torch.nn.Module,
+        parallelism_config: ParallelismConfig,
+        **kwargs,
+    ) -> torch.nn.Module:
+        assert (
+            parallelism_config.tp_size is not None
+            and parallelism_config.tp_size > 1
+        ), (
+            "parallel_config.tp_size must be set and greater than 1 for "
+            "tensor parallelism"
+        )
+
+        device_type = torch.accelerator.current_accelerator().type
+        tp_mesh: DeviceMesh = init_device_mesh(
+            device_type=device_type,
+            mesh_shape=[parallelism_config.tp_size],
+        )
+
+        transformer = self.parallelize_transformer(
+            transformer=transformer,
+            tp_mesh=tp_mesh,
+        )
+        # TODO: Parallelize t5 text encoder via `apply_extra`
+        # abstract method and `extra_parallel_kwargs` ?
+
+        return transformer
+
+    def parallelize_t5(
+        self,
+        text_encoder: nn.Module,
+        tp_mesh: DeviceMesh,
+    ):
+        for i, block in enumerate(text_encoder.encoder.block):
+            block.layer[0].SelfAttention.n_heads //= tp_mesh.size()
+            block.layer[0].SelfAttention.inner_dim //= tp_mesh.size()
+            layer_plan = {
+                "layer.0.SelfAttention.q": ColwiseParallel(),
+                "layer.0.SelfAttention.k": ColwiseParallel(),
+                "layer.0.SelfAttention.v": ColwiseParallel(),
+                "layer.0.SelfAttention.o": RowwiseParallel(),
+                "layer.1.DenseReluDense.wi_0": ColwiseParallel(),
+                "layer.1.DenseReluDense.wi_1": ColwiseParallel(),
+                "layer.1.DenseReluDense.wo": RowwiseParallel(),
+            }
+            if i == 0:
+                layer_plan["layer.0.SelfAttention.relative_attention_bias"] = (
+                    ColwiseParallel()
+                )
+            parallelize_module(
+                module=block,
+                device_mesh=tp_mesh,
+                parallelize_plan=layer_plan,
+            )
+
+        return text_encoder
+
+    def parallelize_transformer(
+        self,
+        transformer: nn.Module,
+        tp_mesh: DeviceMesh,
+    ):
+        for _, block in transformer.transformer_blocks.named_children():
+            block.attn.heads //= tp_mesh.size()
+            layer_plan = {
+                "attn.to_q": ColwiseParallel(),
+                "attn.to_k": ColwiseParallel(),
+                "attn.to_v": ColwiseParallel(),
+                "attn.to_out.0": RowwiseParallel(),
+                "norm1.linear": ColwiseParallel(output_layouts=Replicate()),
+                "ff.net.0.proj": ColwiseParallel(),
+                "ff.net.2": RowwiseParallel(),
+                "attn.add_q_proj": ColwiseParallel(),
+                "attn.add_k_proj": ColwiseParallel(),
+                "attn.add_v_proj": ColwiseParallel(),
+                "attn.to_add_out": RowwiseParallel(),
+                "norm1_context.linear": ColwiseParallel(
+                    output_layouts=Replicate()
+                ),
+                "ff_context.net.0.proj": ColwiseParallel(),
+                "ff_context.net.2": RowwiseParallel(),
+            }
+            parallelize_module(
+                module=block,
+                device_mesh=tp_mesh,
+                parallelize_plan=layer_plan,
+            )
+
+        # NOTE: special handling for FluxSingleTransformerBlock, we have to
+        # rearrange the proj_out weight because it contains both out and down
+        # projection weights in a single matrix.
+        def rearrange_proj_out_weight(
+            single_block: FluxSingleTransformerBlock, tp_group_size
+        ):
+            # rowwise
+            hidden_dim = 3072
+            requires_grad = single_block.proj_out.weight.requires_grad
+            linear2_weight_data = (
+                single_block.proj_out.weight.data.T.detach().clone()
+            )
+            out_weight = linear2_weight_data[:hidden_dim, ...]
+            out_weight = rearrange(
+                out_weight, "(G D) C -> G D C", G=tp_group_size
+            )
+            down_weight = linear2_weight_data.data[hidden_dim:, ...]
+            down_weight = rearrange(
+                down_weight, "(G D) C -> G D C", G=tp_group_size
+            )
+            new_linear2_weight = torch.cat([out_weight, down_weight], dim=1)
+            new_linear2_weight = rearrange(
+                new_linear2_weight, "G D C -> (G D) C"
+            )
+            single_block.proj_out.weight.data.copy_(new_linear2_weight.T)
+            single_block.proj_out.weight.requires_grad_(requires_grad)
+
+        for _, block in transformer.single_transformer_blocks.named_children():
+            rearrange_proj_out_weight(block, tp_mesh.size())
+            block.attn.heads //= tp_mesh.size()
+            layer_plan = {
+                "attn.to_q": ColwiseParallel(),
+                "attn.to_k": ColwiseParallel(),
+                "attn.to_v": ColwiseParallel(),
+                "proj_mlp": ColwiseParallel(),
+                "proj_out": RowwiseParallel(),
+                "norm.linear": ColwiseParallel(output_layouts=Replicate()),
+            }
+            parallelize_module(
+                module=block,
+                device_mesh=tp_mesh,
+                parallelize_plan=layer_plan,
+            )
+        return transformer