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

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,58 @@
+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

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,12 @@
+# NOTE: must import all planner classes to register them
+from .tp_plan_registers import TensorParallelismPlannerRegister
+from .tp_plan_flux import FluxTensorParallelismPlanner
+from .tp_plan_qwen_image import QwenImageTensorParallelismPlanner
+from .tp_plan_wan import WanTensorParallelismPlanner
+
+__all__ = [
+    "TensorParallelismPlannerRegister",
+    "FluxTensorParallelismPlanner",
+    "QwenImageTensorParallelismPlanner",
+    "WanTensorParallelismPlanner",
+]

cache_dit/parallelism/parallel_backend.py

@@ -8,6 +8,8 @@ class ParallelismBackend(Enum):
 
     @classmethod
     def is_supported(cls, backend: "ParallelismBackend") -> bool:
+        if backend in [cls.NATIVE_PYTORCH]:
+            return True
         # Now, only Native_Diffuser backend is supported
         if backend in [cls.NATIVE_DIFFUSER]:
             try:

cache_dit/parallelism/parallel_config.py

@@ -34,7 +34,14 @@ class ParallelismConfig:
             f"Parallel backend {self.backend} is not supported. "
             f"Please make sure the required packages are installed."
         )
-        assert self.tp_size is None, "Tensor parallelism is not supported yet."
+
+        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."
 
     def strify(self, details: bool = False) -> str:
         if details:

cache_dit/parallelism/parallel_interface.py

@@ -24,12 +24,17 @@ def enable_parallelism(
     if parallelism_config.backend == ParallelismBackend.NATIVE_DIFFUSER:
         from cache_dit.parallelism.backends.native_diffusers import (
             maybe_enable_parallelism,
-            native_diffusers_parallelism_available,
         )
 
-        assert (
-            native_diffusers_parallelism_available()
-        ), "Please install diffusers>=0.36.dev0 to use Native_Diffuser backend."
+        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,

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}

@@ -9,7 +9,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 +24,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",
@@ -183,7 +195,11 @@ def quantize_ao(
         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"
@@ -199,10 +215,13 @@ def quantize_ao(
     return module
 
 
-def force_empty_cache():
-    time.sleep(1)
-    gc.collect()
-    torch.cuda.empty_cache()
-    time.sleep(1)
-    gc.collect()
-    torch.cuda.empty_cache()
+def maybe_empty_cache():
+    try:
+        time.sleep(1)
+        gc.collect()
+        torch.cuda.empty_cache()
+        time.sleep(1)
+        gc.collect()
+        torch.cuda.empty_cache()
+    except Exception:
+        pass

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,

cache_dit/utils.py

@@ -13,6 +13,7 @@ from cache_dit.cache_factory import CacheType
 from cache_dit.cache_factory import BlockAdapter
 from cache_dit.cache_factory import BasicCacheConfig
 from cache_dit.cache_factory import CalibratorConfig
+from cache_dit.cache_factory import FakeDiffusionPipeline
 from cache_dit.parallelism import ParallelismConfig
 from cache_dit.logger import init_logger
 
@@ -64,6 +65,7 @@ def summary(
     adapter_or_others: Union[
         BlockAdapter,
         DiffusionPipeline,
+        FakeDiffusionPipeline,
         torch.nn.Module,
     ],
     details: bool = False,
@@ -73,9 +75,15 @@ def summary(
     if adapter_or_others is None:
         return [CacheStats()]
 
+    if isinstance(adapter_or_others, FakeDiffusionPipeline):
+        raise ValueError(
+            "Please pass DiffusionPipeline, BlockAdapter or transfomer, "
+            "not FakeDiffusionPipeline."
+        )
+
     if not isinstance(adapter_or_others, BlockAdapter):
         if not isinstance(adapter_or_others, DiffusionPipeline):
-            transformer = adapter_or_others
+            transformer = adapter_or_others  # transformer-only
             transformer_2 = None
         else:
             transformer = adapter_or_others.transformer
@@ -165,11 +173,18 @@ def strify(
     adapter_or_others: Union[
         BlockAdapter,
         DiffusionPipeline,
+        FakeDiffusionPipeline,
+        torch.nn.Module,
         CacheStats,
         List[CacheStats],
         Dict[str, Any],
     ],
 ) -> str:
+    if isinstance(adapter_or_others, FakeDiffusionPipeline):
+        raise ValueError(
+            "Please pass DiffusionPipeline, BlockAdapter or transfomer, "
+            "not FakeDiffusionPipeline."
+        )
 
     parallelism_config: ParallelismConfig = None
     if isinstance(adapter_or_others, BlockAdapter):
@@ -180,6 +195,10 @@ def strify(
         stats = summary(adapter_or_others, logging=False)[-1]
         cache_options = stats.cache_options
         cached_steps = len(stats.cached_steps)
+    elif isinstance(adapter_or_others, torch.nn.Module):
+        stats = summary(adapter_or_others, logging=False)[-1]
+        cache_options = stats.cache_options
+        cached_steps = len(stats.cached_steps)
     elif isinstance(adapter_or_others, CacheStats):
         stats = adapter_or_others
         cache_options = stats.cache_options
@@ -202,7 +221,8 @@ def strify(
     else:
         raise ValueError(
             "Please set pipe_or_stats param as one of: "
-            "DiffusionPipeline | CacheStats | Dict[str, Any]"
+            "DiffusionPipeline | CacheStats | Dict[str, Any] | List[CacheStats]"
+            " | BlockAdapter | Transformer"
         )
 
     if stats is not None: