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

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

@@ -0,0 +1,153 @@
+from typing import Optional, Union
+
+import torch
+from torch import nn
+from torch.distributed import DeviceMesh, init_device_mesh
+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__)
+
+
+class DistributedRMSNorm(nn.Module):
+    def __init__(
+        self,
+        tp_mesh: DeviceMesh,
+        normalized_shape: Union[int, list[int], torch.Size],
+        eps: Optional[float],
+        elementwise_affine: bool,
+        weight: torch.nn.parameter.Parameter,
+    ):
+        super().__init__()
+        self.tp_mesh = tp_mesh
+        self.elementwise_affine = elementwise_affine
+        self.normalized_shape = normalized_shape
+        self.eps = eps
+        if self.elementwise_affine:
+            assert weight is not None
+        self.weight = weight
+
+    @classmethod
+    def from_rmsnorm(cls, tp_mesh: DeviceMesh, rmsnorm: nn.RMSNorm):
+        if not isinstance(rmsnorm, int):
+            assert len(rmsnorm.normalized_shape) == 1
+
+        if rmsnorm.weight is not None:
+            tp_size = tp_mesh.get_group().size()
+            tp_rank = tp_mesh.get_group().rank()
+            weight = rmsnorm.weight.chunk(tp_size, dim=0)[tp_rank]
+        else:
+            weight = None
+        norm = cls(
+            tp_mesh=tp_mesh,
+            normalized_shape=rmsnorm.normalized_shape,
+            eps=rmsnorm.eps,
+            elementwise_affine=rmsnorm.elementwise_affine,
+            weight=weight,
+        )
+        return norm
+
+    def forward(self, x):
+        if self.elementwise_affine:
+            assert x.shape[-1] == self.weight.shape[0]
+        mean_square = torch.mean(x * x, dim=-1, keepdim=True)
+        torch.distributed.all_reduce(
+            mean_square,
+            op=torch.distributed.ReduceOp.AVG,
+            group=self.tp_mesh.get_group(),
+        )
+        root_mean_square = torch.sqrt(mean_square + self.eps)
+        x_normed = x / root_mean_square
+        if self.elementwise_affine:
+            x_normed = x_normed * self.weight.to(device=x.device)
+        assert x_normed.device.type != "cpu"
+        return x_normed
+
+
+@TensorParallelismPlannerRegister.register("Wan")
+class WanTensorParallelismPlanner(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.blocks.named_children():
+            block.attn1.heads //= tp_mesh.size()
+            block.attn2.heads //= tp_mesh.size()
+            layer_plan = {
+                "attn1.to_q": ColwiseParallel(),
+                "attn1.to_k": ColwiseParallel(),
+                "attn1.to_v": ColwiseParallel(),
+                "attn1.to_out.0": RowwiseParallel(),
+                "attn2.to_q": ColwiseParallel(),
+                "attn2.to_k": ColwiseParallel(),
+                "attn2.to_v": ColwiseParallel(),
+                "attn2.to_out.0": RowwiseParallel(),
+                "ffn.net.0.proj": ColwiseParallel(),
+                "ffn.net.2": RowwiseParallel(),
+            }
+            if getattr(block.attn2, "add_k_proj", None):
+                layer_plan["attn2.add_k_proj"] = ColwiseParallel()
+            if getattr(block.attn2, "add_v_proj", None):
+                layer_plan["attn2.add_v_proj"] = ColwiseParallel()
+            parallelize_module(
+                module=block,
+                device_mesh=tp_mesh,
+                parallelize_plan=layer_plan,
+            )
+
+            block.attn1.norm_q = DistributedRMSNorm.from_rmsnorm(
+                tp_mesh, block.attn1.norm_q
+            )
+            block.attn1.norm_k = DistributedRMSNorm.from_rmsnorm(
+                tp_mesh, block.attn1.norm_k
+            )
+            block.attn2.norm_q = DistributedRMSNorm.from_rmsnorm(
+                tp_mesh, block.attn2.norm_q
+            )
+            block.attn2.norm_k = DistributedRMSNorm.from_rmsnorm(
+                tp_mesh, block.attn2.norm_k
+            )
+            if getattr(block.attn2, "norm_added_k", None):
+                block.attn2.norm_added_k = DistributedRMSNorm.from_rmsnorm(
+                    tp_mesh, block.attn2.norm_added_k
+                )
+        return transformer

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

@@ -0,0 +1,14 @@
+# NOTE: must import all planner classes to register them
+from .tp_plan_flux import FluxTensorParallelismPlanner
+from .tp_plan_kandinsky5 import Kandinsky5TensorParallelismPlanner
+from .tp_plan_qwen_image import QwenImageTensorParallelismPlanner
+from .tp_plan_registers import TensorParallelismPlannerRegister
+from .tp_plan_wan import WanTensorParallelismPlanner
+
+__all__ = [
+    "FluxTensorParallelismPlanner",
+    "Kandinsky5TensorParallelismPlanner",
+    "QwenImageTensorParallelismPlanner",
+    "TensorParallelismPlannerRegister",
+    "WanTensorParallelismPlanner",
+]

cache_dit/parallelism/parallel_backend.py

@@ -0,0 +1,26 @@
+from enum import Enum
+
+
+class ParallelismBackend(Enum):
+    NATIVE_DIFFUSER = "Native_Diffuser"
+    NATIVE_PYTORCH = "Native_PyTorch"
+    NONE = "None"
+
+    @classmethod
+    def is_supported(cls, backend: "ParallelismBackend") -> bool:
+        if backend == cls.NATIVE_PYTORCH:
+            return True
+        elif backend == cls.NATIVE_DIFFUSER:
+            try:
+                from diffusers.models._modeling_parallel import (  # noqa F401
+                    ContextParallelModelPlan,
+                )
+            except ImportError:
+                raise ImportError(
+                    "NATIVE_DIFFUSER parallelism backend requires the latest "
+                    "version of diffusers(>=0.36.dev0). Please install latest "
+                    "version of diffusers from source: \npip3 install "
+                    "git+https://github.com/huggingface/diffusers.git"
+                )
+            return True
+        return False

cache_dit/parallelism/parallel_config.py

@@ -0,0 +1,88 @@
+import dataclasses
+from typing import Optional, Dict, Any
+from cache_dit.parallelism.parallel_backend import ParallelismBackend
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+@dataclasses.dataclass
+class ParallelismConfig:
+    # Parallelism backend, defaults to NATIVE_DIFFUSER
+    backend: ParallelismBackend = ParallelismBackend.NATIVE_DIFFUSER
+    # Context parallelism config
+    # ulysses_size (`int`, *optional*):
+    #     The degree of ulysses parallelism.
+    ulysses_size: int = None
+    # ring_size (`int`, *optional*):
+    #     The degree of ring parallelism.
+    ring_size: int = None
+    # Tensor parallelism config
+    # tp_size (`int`, *optional*):
+    #     The degree of tensor parallelism.
+    tp_size: int = None
+    # parallel_kwargs (`dict`, *optional*):
+    #     Additional kwargs for parallelism backends. For example, for
+    #     NATIVE_DIFFUSER backend, it can include `cp_plan` and
+    #     `attention_backend` arguments for `Context Parallelism`.
+    parallel_kwargs: Optional[Dict[str, Any]] = dataclasses.field(
+        default_factory=dict
+    )
+
+    def __post_init__(self):
+        assert ParallelismBackend.is_supported(self.backend), (
+            f"Parallel backend {self.backend} is not supported. "
+            f"Please make sure the required packages are installed."
+        )
+
+        # Validate the parallelism configuration and auto adjust the backend if needed
+        if self.tp_size is not None and self.tp_size > 1:
+            assert (
+                self.ulysses_size is None or self.ulysses_size == 1
+            ), "Tensor parallelism plus Ulysses parallelism is not supported right now."
+            assert (
+                self.ring_size is None or self.ring_size == 1
+            ), "Tensor parallelism plus Ring parallelism is not supported right now."
+            if self.backend != ParallelismBackend.NATIVE_PYTORCH:
+                logger.warning(
+                    "Tensor parallelism is only supported for NATIVE_PYTORCH backend "
+                    "right now. Force set backend to NATIVE_PYTORCH."
+                )
+                self.backend = ParallelismBackend.NATIVE_PYTORCH
+        elif (
+            self.ulysses_size is not None
+            and self.ulysses_size > 1
+            and self.ring_size is not None
+            and self.ring_size > 1
+        ):
+            raise ValueError(
+                "Ulysses parallelism plus Ring parallelism is not fully supported right now."
+            )
+        else:
+            if (self.ulysses_size is not None and self.ulysses_size > 1) or (
+                self.ring_size is not None and self.ring_size > 1
+            ):
+                if self.backend != ParallelismBackend.NATIVE_DIFFUSER:
+                    logger.warning(
+                        "Ulysses/Ring parallelism is only supported for NATIVE_DIFFUSER "
+                        "backend right now. Force set backend to NATIVE_DIFFUSER."
+                    )
+                    self.backend = ParallelismBackend.NATIVE_DIFFUSER
+
+    def strify(self, details: bool = False) -> str:
+        if details:
+            return (
+                f"ParallelismConfig(backend={self.backend}, "
+                f"ulysses_size={self.ulysses_size}, "
+                f"ring_size={self.ring_size}, "
+                f"tp_size={self.tp_size})"
+            )
+        else:
+            parallel_str = ""
+            if self.ulysses_size is not None:
+                parallel_str += f"Ulysses{self.ulysses_size}"
+            if self.ring_size is not None:
+                parallel_str += f"Ring{self.ring_size}"
+            if self.tp_size is not None:
+                parallel_str += f"TP{self.tp_size}"
+            return parallel_str

cache_dit/parallelism/parallel_interface.py

@@ -0,0 +1,77 @@
+import torch
+from cache_dit.parallelism.parallel_backend import ParallelismBackend
+from cache_dit.parallelism.parallel_config import ParallelismConfig
+from cache_dit.utils import maybe_empty_cache
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+def enable_parallelism(
+    transformer: torch.nn.Module,
+    parallelism_config: ParallelismConfig,
+) -> torch.nn.Module:
+    assert isinstance(transformer, torch.nn.Module), (
+        "transformer must be an instance of torch.nn.Module, "
+        f"but got {type(transformer)}"
+    )
+    if getattr(transformer, "_is_parallelized", False):
+        logger.warning(
+            "The transformer is already parallelized. "
+            "Skipping parallelism enabling."
+        )
+        return transformer
+
+    if parallelism_config.backend == ParallelismBackend.NATIVE_DIFFUSER:
+        from cache_dit.parallelism.backends.native_diffusers import (
+            maybe_enable_parallelism,
+        )
+
+        transformer = maybe_enable_parallelism(
+            transformer,
+            parallelism_config,
+        )
+    elif parallelism_config.backend == ParallelismBackend.NATIVE_PYTORCH:
+        from cache_dit.parallelism.backends.native_pytorch import (
+            maybe_enable_parallelism,
+        )
+
+        transformer = maybe_enable_parallelism(
+            transformer,
+            parallelism_config,
+        )
+    else:
+        raise ValueError(
+            f"Parallel backend {parallelism_config.backend} is not supported yet."
+        )
+
+    transformer._is_parallelized = True  # type: ignore[attr-defined]
+    # Use `parallelism` not `parallel` to avoid name conflict with diffusers.
+    transformer._parallelism_config = parallelism_config  # type: ignore[attr-defined]
+    logger.info(
+        f"Enabled parallelism: {parallelism_config.strify(True)}, "
+        f"transformer id:{id(transformer)}"
+    )
+
+    # NOTE: Workaround for potential memory peak issue after parallelism
+    # enabling, specially for tensor parallelism in native pytorch backend.
+    maybe_empty_cache()
+
+    return transformer
+
+
+def remove_parallelism_stats(
+    transformer: torch.nn.Module,
+) -> torch.nn.Module:
+    if not getattr(transformer, "_is_parallelized", False):
+        logger.warning(
+            "The transformer is not parallelized. "
+            "Skipping removing parallelism."
+        )
+        return transformer
+
+    if hasattr(transformer, "_is_parallelized"):
+        del transformer._is_parallelized  # type: ignore[attr-defined]
+    if hasattr(transformer, "_parallelism_config"):
+        del transformer._parallelism_config  # type: ignore[attr-defined]
+    return transformer

cache_dit/quantize/__init__.py

@@ -1 +1,8 @@
+try:
+    import torchao
+except ImportError:
+    raise ImportError(
+        "Quantization functionality requires the 'quantization' extra dependencies. "
+        "Install with: pip install cache-dit[quantization]"
+    )
 from cache_dit.quantize.quantize_interface import quantize

cache_dit/quantize/backends/__init__.py

@@ -0,0 +1 @@
+from .torchao import quantize_ao

cache_dit/quantize/backends/torchao/__init__.py

@@ -0,0 +1 @@
+from .quantize_ao import quantize_ao

cache_dit/quantize/{quantize_ao.py → backends/torchao/quantize_ao.py}

@@ -1,7 +1,6 @@
-import gc
-import time
 import torch
 from typing import Callable, Optional, List
+from cache_dit.utils import maybe_empty_cache
 from cache_dit.logger import init_logger
 
 logger = init_logger(__name__)
@@ -9,7 +8,7 @@ logger = init_logger(__name__)
 
 def quantize_ao(
     module: torch.nn.Module,
-    quant_type: str = "fp8_w8a8_dq",
+    quant_type: str = "float8_weight_only",
     exclude_layers: List[str] = [
         "embedder",
         "embed",
@@ -24,6 +23,18 @@ def quantize_ao(
     # set `exclude_layers` as `[]` if you don't want this behavior.
     assert isinstance(module, torch.nn.Module)
 
+    alias_map = {
+        "float8": "fp8_w8a8_dq",
+        "float8_weight_only": "fp8_w8a16_wo",
+        "int8": "int8_w8a8_dq",
+        "int8_weight_only": "int8_w8a16_wo",
+        "int4": "int4_w4a8_dq",
+        "int4_w4a4": "int4_w4a4_dq",
+        "int4_weight_only": "int4_w4a16_wo",
+    }
+    if quant_type.lower() in alias_map:
+        quant_type = alias_map[quant_type.lower()]
+
     quant_type = quant_type.lower()
     assert quant_type in (
         "fp8_w8a8_dq",
@@ -80,11 +91,11 @@ def quantize_ao(
 
         return False
 
-    def _quantization_fn():
+    def _quant_config():
         try:
             if quant_type == "fp8_w8a8_dq":
                 from torchao.quantization import (
-                    float8_dynamic_activation_float8_weight,
+                    Float8DynamicActivationFloat8WeightConfig,
                     PerTensor,
                     PerRow,
                 )
@@ -92,7 +103,7 @@ def quantize_ao(
                 if per_row:  # Ensure bfloat16
                     module.to(torch.bfloat16)
 
-                quantization_fn = float8_dynamic_activation_float8_weight(
+                quant_config = Float8DynamicActivationFloat8WeightConfig(
                     weight_dtype=kwargs.get(
                         "weight_dtype",
                         torch.float8_e4m3fn,
@@ -109,9 +120,9 @@ def quantize_ao(
                 )
 
             elif quant_type == "fp8_w8a16_wo":
-                from torchao.quantization import float8_weight_only
+                from torchao.quantization import Float8WeightOnlyConfig
 
-                quantization_fn = float8_weight_only(
+                quant_config = Float8WeightOnlyConfig(
                     weight_dtype=kwargs.get(
                         "weight_dtype",
                         torch.float8_e4m3fn,
@@ -120,39 +131,43 @@ def quantize_ao(
 
             elif quant_type == "int8_w8a8_dq":
                 from torchao.quantization import (
-                    int8_dynamic_activation_int8_weight,
+                    Int8DynamicActivationInt8WeightConfig,
                 )
 
-                quantization_fn = int8_dynamic_activation_int8_weight()
+                quant_config = Int8DynamicActivationInt8WeightConfig()
 
             elif quant_type == "int8_w8a16_wo":
-                from torchao.quantization import int8_weight_only
 
-                quantization_fn = int8_weight_only(
+                from torchao.quantization import Int8WeightOnlyConfig
+
+                quant_config = Int8WeightOnlyConfig(
                     # group_size is None -> per_channel, else per group
                     group_size=kwargs.get("group_size", None),
                 )
 
             elif quant_type == "int4_w4a8_dq":
+
                 from torchao.quantization import (
-                    int8_dynamic_activation_int4_weight,
+                    Int8DynamicActivationInt4WeightConfig,
                 )
 
-                quantization_fn = int8_dynamic_activation_int4_weight(
+                quant_config = Int8DynamicActivationInt4WeightConfig(
                     group_size=kwargs.get("group_size", 32),
                 )
 
             elif quant_type == "int4_w4a4_dq":
+
                 from torchao.quantization import (
-                    int4_dynamic_activation_int4_weight,
+                    Int4DynamicActivationInt4WeightConfig,
                 )
 
-                quantization_fn = int4_dynamic_activation_int4_weight()
+                quant_config = Int4DynamicActivationInt4WeightConfig()
 
             elif quant_type == "int4_w4a16_wo":
-                from torchao.quantization import int4_weight_only
 
-                quantization_fn = int4_weight_only(
+                from torchao.quantization import Int4WeightOnlyConfig
+
+                quant_config = Int4WeightOnlyConfig(
                     group_size=kwargs.get("group_size", 32),
                 )
 
@@ -168,33 +183,32 @@ def quantize_ao(
             )
             raise e
 
-        return quantization_fn
+        return quant_config
 
     from torchao.quantization import quantize_
 
     quantize_(
         module,
-        _quantization_fn(),
+        _quant_config(),
         filter_fn=_filter_fn if filter_fn is None else filter_fn,
         device=kwargs.get("device", None),
     )
 
-    force_empty_cache()
+    maybe_empty_cache()
+
+    alias_map_rev = {v: k for k, v in alias_map.items()}
+    if quant_type in alias_map_rev:
+        quant_type = alias_map_rev[quant_type]
 
     logger.info(
+        f"Quantized        Module: {module.__class__.__name__:>5}\n"
         f"Quantized        Method: {quant_type:>5}\n"
         f"Quantized Linear Layers: {num_quant_linear:>5}\n"
         f"Skipped   Linear Layers: {num_skip_linear:>5}\n"
         f"Total     Linear Layers: {num_linear_layers:>5}\n"
         f"Total     (all)  Layers: {num_layers:>5}"
     )
-    return module
 
-
-def force_empty_cache():
-    time.sleep(1)
-    gc.collect()
-    torch.cuda.empty_cache()
-    time.sleep(1)
-    gc.collect()
-    torch.cuda.empty_cache()
+    module._quantize_type = quant_type
+    module._is_quantized = True
+    return module

cache_dit/quantize/quantize_interface.py

@@ -7,37 +7,24 @@ logger = init_logger(__name__)
 
 def quantize(
     module: torch.nn.Module,
-    quant_type: str = "fp8_w8a8_dq",
+    quant_type: str = "float8_weight_only",
     backend: str = "ao",
     exclude_layers: List[str] = [
         "embedder",
         "embed",
     ],
     filter_fn: Optional[Callable] = None,
-    # only for fp8_w8a8_dq
-    per_row: bool = True,
     **kwargs,
 ) -> torch.nn.Module:
     assert isinstance(module, torch.nn.Module)
 
     if backend.lower() in ("ao", "torchao"):
-        from cache_dit.quantize.quantize_ao import quantize_ao
-
-        quant_type = quant_type.lower()
-        assert quant_type in (
-            "fp8_w8a8_dq",
-            "fp8_w8a16_wo",
-            "int8_w8a8_dq",
-            "int8_w8a16_wo",
-            "int4_w4a8_dq",
-            "int4_w4a4_dq",
-            "int4_w4a16_wo",
-        ), f"{quant_type} is not supported for torchao backend now!"
+        from cache_dit.quantize.backends.torchao import quantize_ao
 
         return quantize_ao(
             module,
             quant_type=quant_type,
-            per_row=per_row,
+            per_row=kwargs.pop("per_row", True),
             exclude_layers=exclude_layers,
             filter_fn=filter_fn,
             **kwargs,