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

cache_dit/cache_factory/forward_pattern.py

@@ -20,33 +20,33 @@ class ForwardPattern(Enum):
 
     Pattern_0 = (
         True,  # Return_H_First
-        False, # Return_H_Only
-        False, # Forward_H_only
+        False,  # Return_H_Only
+        False,  # Forward_H_only
         ("hidden_states", "encoder_hidden_states"),  # In
         ("hidden_states", "encoder_hidden_states"),  # Out
         True,  # Supported
     )
 
     Pattern_1 = (
-        False, # Return_H_First
-        False, # Return_H_Only
-        False, # Forward_H_only
+        False,  # Return_H_First
+        False,  # Return_H_Only
+        False,  # Forward_H_only
         ("hidden_states", "encoder_hidden_states"),  # In
         ("encoder_hidden_states", "hidden_states"),  # Out
         True,  # Supported
     )
 
     Pattern_2 = (
-        False, # Return_H_First
+        False,  # Return_H_First
         True,  # Return_H_Only
-        False, # Forward_H_only
+        False,  # Forward_H_only
         ("hidden_states", "encoder_hidden_states"),  # In
-        ("hidden_states",),                          # Out
+        ("hidden_states",),  # Out
         True,  # Supported
     )
 
     Pattern_3 = (
-        False, # Return_H_First
+        False,  # Return_H_First
         True,  # Return_H_Only
         True,  # Forward_H_only
         ("hidden_states",),  # In
@@ -56,18 +56,18 @@ class ForwardPattern(Enum):
 
     Pattern_4 = (
         True,  # Return_H_First
-        False, # Return_H_Only
+        False,  # Return_H_Only
         True,  # Forward_H_only
-        ("hidden_states",),                          # In
+        ("hidden_states",),  # In
         ("hidden_states", "encoder_hidden_states"),  # Out
         True,  # Supported
     )
 
     Pattern_5 = (
-        False, # Return_H_First
-        False, # Return_H_Only
+        False,  # Return_H_First
+        False,  # Return_H_Only
         True,  # Forward_H_only
-        ("hidden_states",),                          # In
+        ("hidden_states",),  # In
         ("encoder_hidden_states", "hidden_states"),  # Out
         True,  # Supported
     )

cache_dit/parallelism/backends/native_diffusers/__init__.py

@@ -1,6 +1,3 @@
 from cache_dit.parallelism.backends.native_diffusers.parallel_difffusers import (
     maybe_enable_parallelism,
 )
-from cache_dit.parallelism.backends.native_diffusers.parallel_difffusers import (
-    native_diffusers_parallelism_available,
-)

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

@@ -0,0 +1,95 @@
+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 .cp_planners import *
+
+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_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
+                )
+            else:
+                raise ValueError(
+                    f"{transformer.__class__.__name__} does not support context parallelism."
+                )
+
+    return transformer

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

@@ -0,0 +1,74 @@
+import torch
+import logging
+from abc import abstractmethod
+from typing import Optional
+from diffusers.models.modeling_utils import ModelMixin
+
+try:
+    from diffusers.models._modeling_parallel import (
+        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 cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+__all__ = [
+    "ContextParallelismPlanner",
+    "ContextParallelismPlannerRegister",
+]
+
+
+class ContextParallelismPlanner:
+    @abstractmethod
+    def apply(
+        self,
+        # NOTE: Keep this kwarg for future extensions
+        transformer: Optional[torch.nn.Module | ModelMixin] = None,
+        **kwargs,
+    ) -> ContextParallelModelPlan:
+        # NOTE: This method should only return the CP plan dictionary.
+        raise NotImplementedError(
+            "apply method must be implemented by subclasses"
+        )
+
+
+class ContextParallelismPlannerRegister:
+    _cp_planner_registry: dict[str, ContextParallelismPlanner] = {}
+
+    @classmethod
+    def register(cls, name: str):
+        def decorator(planner_cls: type[ContextParallelismPlanner]):
+            assert (
+                name not in cls._cp_planner_registry
+            ), f"ContextParallelismPlanner with name {name} is already registered."
+            if logger.isEnabledFor(logging.DEBUG):
+                logger.debug(f"Registering ContextParallelismPlanner: {name}")
+            cls._cp_planner_registry[name] = planner_cls
+            return planner_cls
+
+        return decorator
+
+    @classmethod
+    def get_planner(
+        cls, transformer: str | torch.nn.Module | ModelMixin
+    ) -> type[ContextParallelismPlanner]:
+        if isinstance(transformer, (torch.nn.Module, ModelMixin)):
+            name = transformer.__class__.__name__
+        else:
+            name = transformer
+        planner_cls = None
+        for planner_name in cls._cp_planner_registry:
+            if name.startswith(planner_name):
+                planner_cls = cls._cp_planner_registry.get(planner_name)
+                break
+        if planner_cls is None:
+            raise ValueError(f"No planner registered under name: {name}")
+        return planner_cls

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,42 +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:
-            if hasattr(transformer, "enable_parallelism"):
-                if hasattr(transformer, "set_attention_backend"):
-                    # Now only _native_cudnn is supported for parallelism
-                    # issue: https://github.com/huggingface/diffusers/pull/12443
-                    transformer.set_attention_backend("_native_cudnn")
-                    logger.warning(
-                        "Set attention backend to _native_cudnn for parallelism because of "
-                        "the issue: https://github.com/huggingface/diffusers/pull/12443"
-                    )
-                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