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

cache_dit/parallelism/backends/native_diffusers/parallel_difffusers.py

@@ -0,0 +1,49 @@
+import torch
+
+from typing import Optional
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+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(
+    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)}"
+    )
+
+    assert parallelism_config.backend == ParallelismBackend.NATIVE_DIFFUSER, (
+        f"parallelism backend must be {ParallelismBackend.NATIVE_DIFFUSER}, "
+        f"but got {parallelism_config.backend}"
+    )
+
+    if (
+        parallelism_config.ulysses_size is not None
+        or parallelism_config.ring_size is not None
+    ):
+        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,6 @@
+from cache_dit.parallelism.backends.native_pytorch.tensor_parallelism import (
+    TensorParallelismPlannerRegister,
+)
+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,171 @@
+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.logger import init_logger
+from cache_dit.parallelism.parallel_config import ParallelismConfig
+
+from .tp_plan_registers import (
+    TensorParallelismPlanner,
+    TensorParallelismPlannerRegister,
+)
+
+logger = init_logger(__name__)
+
+
+@TensorParallelismPlannerRegister.register("Chroma")
+@TensorParallelismPlannerRegister.register("HunyuanImage")
+@TensorParallelismPlannerRegister.register("HunyuanVideo")
+@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(),
+                "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(),
+                "ff_context.net.0.proj": ColwiseParallel(),
+                "ff_context.net.2": RowwiseParallel(),
+            }
+
+            if getattr(block.norm1, "linear", None) is not None:
+                layer_plan["norm1.linear"] = ColwiseParallel(
+                    output_layouts=Replicate()
+                )
+            if getattr(block.norm1_context, "linear", None) is not None:
+                layer_plan["norm1_context.linear"] = ColwiseParallel(
+                    output_layouts=Replicate()
+                )
+            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 = single_block.attn.to_q.weight.shape[0]
+            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(),
+            }
+            if getattr(block.norm, "linear", None) is not None:
+                layer_plan["norm.linear"] = ColwiseParallel(
+                    output_layouts=Replicate()
+                )
+            parallelize_module(
+                module=block,
+                device_mesh=tp_mesh,
+                parallelize_plan=layer_plan,
+            )
+        return transformer

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

@@ -0,0 +1,79 @@
+import torch
+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.logger import init_logger
+from cache_dit.parallelism.parallel_config import ParallelismConfig
+
+from .tp_plan_registers import (
+    TensorParallelismPlanner,
+    TensorParallelismPlannerRegister,
+)
+
+logger = init_logger(__name__)
+
+
+@TensorParallelismPlannerRegister.register("Kandinsky5")
+class Kandinsky5TensorParallelismPlanner(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,
+        )
+
+        return transformer
+
+    def parallelize_transformer(
+        self,
+        transformer: nn.Module,
+        tp_mesh: DeviceMesh,
+    ):
+        for _, block in transformer.visual_transformer_blocks.named_children():
+            block.self_attention.num_heads //= tp_mesh.size()
+            block.cross_attention.num_heads //= tp_mesh.size()
+            layer_plan = {
+                "self_attention.to_query": ColwiseParallel(),
+                "self_attention.to_key": ColwiseParallel(),
+                "self_attention.to_value": ColwiseParallel(),
+                "self_attention.out_layer": RowwiseParallel(),
+                "cross_attention.to_query": ColwiseParallel(),
+                "cross_attention.to_key": ColwiseParallel(),
+                "cross_attention.to_value": ColwiseParallel(),
+                "cross_attention.out_layer": RowwiseParallel(),
+                "visual_modulation.out_layer": ColwiseParallel(
+                    output_layouts=Replicate()
+                ),
+                "feed_forward.in_layer": ColwiseParallel(),
+                "feed_forward.out_layer": RowwiseParallel(),
+            }
+            parallelize_module(
+                module=block,
+                device_mesh=tp_mesh,
+                parallelize_plan=layer_plan,
+            )
+        return transformer

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

@@ -0,0 +1,78 @@
+import torch
+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("QwenImage")
+class QwenImageTensorParallelismPlanner(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,
+        )
+
+        return transformer
+
+    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(),
+                "img_mod.1": ColwiseParallel(output_layouts=Replicate()),
+                "img_mlp.net.0.proj": ColwiseParallel(),
+                "img_mlp.net.2": RowwiseParallel(),
+                "attn.add_q_proj": ColwiseParallel(),
+                "attn.add_k_proj": ColwiseParallel(),
+                "attn.add_v_proj": ColwiseParallel(),
+                "attn.to_add_out": RowwiseParallel(),
+                "txt_mod.1": ColwiseParallel(output_layouts=Replicate()),
+                "txt_mlp.net.0.proj": ColwiseParallel(),
+                "txt_mlp.net.2": RowwiseParallel(),
+            }
+            parallelize_module(
+                module=block,
+                device_mesh=tp_mesh,
+                parallelize_plan=layer_plan,
+            )
+        return transformer

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

@@ -0,0 +1,65 @@
+import torch
+import logging
+from abc import abstractmethod
+from typing import Dict
+from cache_dit.parallelism.parallel_config import ParallelismConfig
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class TensorParallelismPlanner:
+    # TODO: add `apply_extra` abstract method for extra
+    # parallelism kwargs handling
+
+    @abstractmethod
+    def apply(
+        self,
+        transformer: torch.nn.Module,
+        parallelism_config: ParallelismConfig,
+        **kwargs,
+    ) -> torch.nn.Module:
+        raise NotImplementedError(
+            "apply method must be implemented by subclasses"
+        )
+
+
+class TensorParallelismPlannerRegister:
+    _tp_planner_registry: Dict[str, TensorParallelismPlanner] = {}
+
+    @classmethod
+    def register(cls, name: str):
+        def decorator(planner_cls: type[TensorParallelismPlanner]):
+            assert (
+                name not in cls._tp_planner_registry
+            ), f"TensorParallelismPlanner with name {name} is already registered."
+            if logger.isEnabledFor(logging.DEBUG):
+                logger.debug(f"Registering TensorParallelismPlanner: {name}")
+            cls._tp_planner_registry[name] = planner_cls
+            return planner_cls
+
+        return decorator
+
+    @classmethod
+    def get_planner(
+        cls, transformer: str | torch.nn.Module
+    ) -> type[TensorParallelismPlanner]:
+        if isinstance(transformer, torch.nn.Module):
+            name = transformer.__class__.__name__
+        else:
+            name = transformer
+        planner_cls = None
+        for planner_name in cls._tp_planner_registry:
+            if name.startswith(planner_name):
+                planner_cls = cls._tp_planner_registry.get(planner_name)
+                break
+        if planner_cls is None:
+            raise ValueError(f"No planner registered under name: {name}")
+        return planner_cls
+
+    @classmethod
+    def supported_planners(
+        cls,
+    ) -> tuple[int, list[str]]:
+        val_planners = cls._tp_planner_registry.keys()
+        return len(val_planners), [p for p in val_planners]