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

cache_dit/{cache_factory → caching}/cache_contexts/calibrators/taylorseer.py

@@ -1,7 +1,9 @@
+# The TaylorSeerState codebase is adapted from FBCache. Over time its codebase
+# diverged a lot, and TaylorSeerState API is no longer compatible with FBCache.
 import math
 import torch
 from typing import List, Dict
-from cache_dit.cache_factory.cache_contexts.calibrators.base import (
+from cache_dit.caching.cache_contexts.calibrators.base import (
     CalibratorBase,
 )
 
@@ -10,13 +12,12 @@ from cache_dit.logger import init_logger
 logger = init_logger(__name__)
 
 
-class TaylorSeerCalibrator(CalibratorBase):
+class TaylorSeerState:
     def __init__(
         self,
         n_derivatives=1,
         max_warmup_steps=1,
         skip_interval_steps=1,
-        **kwargs,
     ):
         self.n_derivatives = n_derivatives
         self.order = n_derivatives + 1
@@ -28,9 +29,8 @@ class TaylorSeerCalibrator(CalibratorBase):
             "dY_prev": [None] * self.order,
             "dY_current": [None] * self.order,
         }
-        self.reset_cache()
 
-    def reset_cache(self):  # NEED
+    def reset(self):
         self.state: Dict[str, List[torch.Tensor]] = {
             "dY_prev": [None] * self.order,
             "dY_current": [None] * self.order,
@@ -38,6 +38,9 @@ class TaylorSeerCalibrator(CalibratorBase):
         self.current_step = -1
         self.last_non_approximated_step = -1
 
+    def mark_step_begin(self):  # NEED
+        self.current_step += 1
+
     def should_compute(self, step=None):
         step = self.current_step if step is None else step
         if (
@@ -56,7 +59,7 @@ class TaylorSeerCalibrator(CalibratorBase):
         window = self.current_step - self.last_non_approximated_step
         if self.state["dY_prev"][0] is not None:
             if dY_current[0].shape != self.state["dY_prev"][0].shape:
-                self.reset_cache()
+                self.reset()
 
         for i in range(self.n_derivatives):
             if self.state["dY_prev"][i] is not None and self.current_step > 1:
@@ -77,9 +80,6 @@ class TaylorSeerCalibrator(CalibratorBase):
                 break
         return output
 
-    def mark_step_begin(self):  # NEED
-        self.current_step += 1
-
     def update(self, Y: torch.Tensor):  # NEED
         # Directly call this method will ingnore the warmup
         # policy and force full computation.
@@ -106,5 +106,77 @@ class TaylorSeerCalibrator(CalibratorBase):
         else:
             return self.approximate()
 
+
+class TaylorSeerCalibrator(CalibratorBase):
+    def __init__(
+        self,
+        n_derivatives=1,
+        max_warmup_steps=1,
+        skip_interval_steps=1,
+        **kwargs,
+    ):
+        self.n_derivatives = n_derivatives
+        self.max_warmup_steps = max_warmup_steps
+        self.skip_interval_steps = skip_interval_steps
+        self.states: Dict[str, TaylorSeerState] = {}
+        self.reset_cache()
+
+    def reset_cache(self):  # NEED
+        if self.states:
+            for state in self.states.values():
+                state.reset()
+
+    def maybe_init_state(
+        self,
+        name: str = "default",
+    ):
+        if name not in self.states:
+            self.states[name] = TaylorSeerState(
+                n_derivatives=self.n_derivatives,
+                max_warmup_steps=self.max_warmup_steps,
+                skip_interval_steps=self.skip_interval_steps,
+            )
+
+    def mark_step_begin(self, *args, **kwargs):
+        if self.states:
+            for state in self.states.values():
+                state.mark_step_begin()
+
+    def derivative(
+        self,
+        Y: torch.Tensor,
+        name: str = "default",
+    ) -> List[torch.Tensor]:
+        self.maybe_init_state(name)
+        state = self.states[name]
+        state.derivative(Y)
+        return state.state["dY_current"]
+
+    def approximate(
+        self,
+        name: str = "default",
+    ) -> torch.Tensor:  # NEED
+        assert name in self.states, f"State '{name}' not found."
+        state = self.states[name]
+        return state.approximate()
+
+    def update(
+        self,
+        Y: torch.Tensor,
+        name: str = "default",
+    ):  # NEED
+        self.maybe_init_state(name)
+        state = self.states[name]
+        state.update(Y)
+
+    def step(
+        self,
+        Y: torch.Tensor,
+        name: str = "default",
+    ):
+        self.maybe_init_state(name)
+        state = self.states[name]
+        return state.step(Y)
+
     def __repr__(self):
         return f"TaylorSeerCalibrator_O({self.n_derivatives})"

cache_dit/caching/cache_contexts/context_manager.py

@@ -0,0 +1,36 @@
+from cache_dit.caching.cache_types import CacheType
+from cache_dit.caching.cache_contexts.cache_manager import (
+    CachedContextManager,
+)
+from cache_dit.caching.cache_contexts.prune_manager import (
+    PrunedContextManager,
+)
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class ContextManager:
+    _supported_managers = (
+        CachedContextManager,
+        PrunedContextManager,
+    )
+
+    def __new__(
+        cls,
+        cache_type: CacheType,
+        name: str = "default",
+        persistent_context: bool = False,
+    ) -> CachedContextManager | PrunedContextManager:
+        if cache_type == CacheType.DBCache:
+            return CachedContextManager(
+                name=name,
+                persistent_context=persistent_context,
+            )
+        elif cache_type == CacheType.DBPrune:
+            return PrunedContextManager(
+                name=name,
+                persistent_context=persistent_context,
+            )
+        else:
+            raise ValueError(f"Unsupported cache_type: {cache_type}.")

cache_dit/caching/cache_contexts/prune_config.py

@@ -0,0 +1,63 @@
+import dataclasses
+from typing import List
+from cache_dit.caching.cache_types import CacheType
+from cache_dit.caching.cache_contexts.cache_config import (
+    BasicCacheConfig,
+)
+
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+@dataclasses.dataclass
+class DBPruneConfig(BasicCacheConfig):
+    # Dyanamic Block Prune specific configurations
+    cache_type: CacheType = CacheType.DBPrune  # DBPrune
+
+    # enable_dynamic_prune_threshold (`bool`, *required*, defaults to False):
+    #     Whether to enable the dynamic prune threshold or not. If True, we will
+    #     compute the dynamic prune threshold based on the mean of the residual
+    #     diffs of the previous computed or pruned blocks.
+    #     But, also limit mean_diff to be at least 2x the residual_diff_threshold
+    #     to avoid too aggressive pruning.
+    enable_dynamic_prune_threshold: bool = False
+    # max_dynamic_prune_threshold (`float`, *optional*, defaults to None):
+    #     The max dynamic prune threshold, if not None, the dynamic prune threshold
+    #     will not exceed this value. If None, we will limit it to be at least 2x
+    #     the residual_diff_threshold to avoid too aggressive pruning.
+    max_dynamic_prune_threshold: float = None
+    # dynamic_prune_threshold_relax_ratio (`float`, *optional*, defaults to 1.25):
+    #     The relax ratio for dynamic prune threshold, the dynamic prune threshold
+    #     will be set as:
+    #         dynamic_prune_threshold = mean_diff * dynamic_prune_threshold_relax_ratio
+    #     to avoid too aggressive pruning.
+    #     The default value is 1.25, which means the dynamic prune threshold will
+    #     be 1.25 times the mean of the residual diffs of the previous computed
+    #     or pruned blocks.
+    #     Users can tune this value to achieve a better trade-off between speedup
+    #     and precision. A higher value leads to more aggressive pruning
+    #     and faster speedup, but may also lead to lower precision.
+    dynamic_prune_threshold_relax_ratio: float = 1.25
+    # non_prune_block_ids (`List[int]`, *optional*, defaults to []):
+    #     The list of block ids that will not be pruned, even if their residual
+    #     diffs are below the prune threshold. This can be useful for the first
+    #     few blocks, which are usually more important for the model performance.
+    non_prune_block_ids: List[int] = dataclasses.field(default_factory=list)
+    # force_reduce_calibrator_vram (`bool`, *optional*, defaults to True):
+    #     Whether to force reduce the VRAM usage of the calibrator for Dynamic Block
+    #     Prune. If True, we will set the downsample_factor of the extra_cache_config
+    #     to at least 2 to reduce the VRAM usage of the calibrator.
+    force_reduce_calibrator_vram: bool = False
+
+    def strify(self) -> str:
+        return (
+            f"{self.cache_type}_"
+            f"F{self.Fn_compute_blocks}"
+            f"B{self.Bn_compute_blocks}_"
+            f"W{self.max_warmup_steps}"
+            f"I{self.warmup_interval}"
+            f"M{max(0, self.max_cached_steps)}"
+            f"MC{max(0, self.max_continuous_cached_steps)}_"
+            f"R{self.residual_diff_threshold}"
+        )

cache_dit/caching/cache_contexts/prune_context.py

@@ -0,0 +1,155 @@
+import torch
+import logging
+import dataclasses
+from typing import List
+from cache_dit.caching.cache_types import CacheType
+from cache_dit.caching.cache_contexts.prune_config import (
+    DBPruneConfig,
+)
+from cache_dit.caching.cache_contexts.cache_context import (
+    CachedContext,
+)
+from cache_dit.logger import init_logger
+
+
+logger = init_logger(__name__)
+
+
+@dataclasses.dataclass
+class PrunedContext(CachedContext):
+    # Overwrite the cache_config type for PrunedContext
+    cache_config: DBPruneConfig = dataclasses.field(
+        default_factory=DBPruneConfig,
+    )
+    # Specially for Dynamic Block Prune
+    pruned_blocks: List[int] = dataclasses.field(default_factory=list)
+    actual_blocks: List[int] = dataclasses.field(default_factory=list)
+    cfg_pruned_blocks: List[int] = dataclasses.field(default_factory=list)
+    cfg_actual_blocks: List[int] = dataclasses.field(default_factory=list)
+
+    def __post_init__(self):
+        super().__post_init__()
+        if not isinstance(self.cache_config, DBPruneConfig):
+            raise ValueError(
+                "PrunedContext only supports DBPruneConfig as cache_config."
+            )
+
+        if self.cache_config.cache_type == CacheType.DBPrune:
+            if (
+                self.calibrator_config is not None
+                and self.cache_config.force_reduce_calibrator_vram
+            ):
+                # May reduce VRAM usage for Dynamic Block Prune
+                self.extra_cache_config.downsample_factor = max(
+                    4, self.extra_cache_config.downsample_factor
+                )
+
+    def get_residual_diff_threshold(self):
+        # Overwite this func for Dynamic Block Prune
+        residual_diff_threshold = self.cache_config.residual_diff_threshold
+        if isinstance(residual_diff_threshold, torch.Tensor):
+            residual_diff_threshold = residual_diff_threshold.item()
+        if self.cache_config.enable_dynamic_prune_threshold:
+            # Compute the dynamic prune threshold based on the mean of the
+            # residual diffs of the previous computed or pruned blocks.
+            step = str(self.get_current_step())
+            if int(step) >= 0 and str(step) in self.residual_diffs:
+                assert isinstance(self.residual_diffs[step], list)
+                # Use all the recorded diffs for this step
+                # NOTE: Should we only use the last 5 diffs?
+                diffs = self.residual_diffs[step][:5]
+                diffs = [d for d in diffs if d > 0.0]
+                if diffs:
+                    mean_diff = sum(diffs) / len(diffs)
+                    relaxed_diff = (
+                        mean_diff
+                        * self.cache_config.dynamic_prune_threshold_relax_ratio
+                    )
+                    if self.cache_config.max_dynamic_prune_threshold is None:
+                        max_dynamic_prune_threshold = (
+                            2 * residual_diff_threshold
+                        )
+                    else:
+                        max_dynamic_prune_threshold = (
+                            self.cache_config.max_dynamic_prune_threshold
+                        )
+                    if relaxed_diff < max_dynamic_prune_threshold:
+                        # If the mean diff is less than twice the threshold,
+                        # we can use it as the dynamic prune threshold.
+                        residual_diff_threshold = (
+                            relaxed_diff
+                            if relaxed_diff > residual_diff_threshold
+                            else residual_diff_threshold
+                        )
+                    if logger.isEnabledFor(logging.DEBUG):
+                        logger.debug(
+                            f"Dynamic prune threshold for step {step}: "
+                            f"{residual_diff_threshold:.6f}"
+                        )
+        return residual_diff_threshold
+
+    def mark_step_begin(self):
+        # Overwite this func for Dynamic Block Prune
+        super().mark_step_begin()
+        # Reset pruned_blocks and actual_blocks at the beginning
+        # of each transformer step.
+        if self.get_current_transformer_step() == 0:
+            self.pruned_blocks.clear()
+            self.actual_blocks.clear()
+
+    def add_residual_diff(self, diff: float | torch.Tensor):
+        # Overwite this func for Dynamic Block Prune
+        if isinstance(diff, torch.Tensor):
+            diff = diff.item()
+        # step: executed_steps - 1, not transformer_steps - 1
+        step = str(self.get_current_step())
+        # For Dynamic Block Prune, we will record all the diffs for this step
+        # Only add the diff if it is not already recorded for this step
+        if not self.is_separate_cfg_step():
+            if step not in self.residual_diffs:
+                self.residual_diffs[step] = []
+            self.residual_diffs[step].append(diff)
+        else:
+            if step not in self.cfg_residual_diffs:
+                self.cfg_residual_diffs[step] = []
+            self.cfg_residual_diffs[step].append(diff)
+
+    def add_pruned_step(self):
+        curr_cached_step = self.get_current_step()
+        # Avoid adding the same step multiple times
+        if not self.is_separate_cfg_step():
+            if curr_cached_step not in self.cached_steps:
+                self.add_cached_step()
+        else:
+            if curr_cached_step not in self.cfg_cached_steps:
+                self.add_cached_step()
+
+    def add_pruned_block(self, num_blocks):
+        if not self.is_separate_cfg_step():
+            self.pruned_blocks.append(num_blocks)
+        else:
+            self.cfg_pruned_blocks.append(num_blocks)
+
+    def add_actual_block(self, num_blocks):
+        if not self.is_separate_cfg_step():
+            self.actual_blocks.append(num_blocks)
+        else:
+            self.cfg_actual_blocks.append(num_blocks)
+
+    def get_pruned_blocks(self):
+        return self.pruned_blocks.copy()
+
+    def get_cfg_pruned_blocks(self):
+        return self.cfg_pruned_blocks.copy()
+
+    def get_actual_blocks(self):
+        return self.actual_blocks.copy()
+
+    def get_cfg_actual_blocks(self):
+        return self.cfg_actual_blocks.copy()
+
+    def get_pruned_steps(self):
+        return self.get_cached_steps()
+
+    def get_cfg_pruned_steps(self):
+        return self.get_cfg_cached_steps()

cache_dit/caching/cache_contexts/prune_manager.py

@@ -0,0 +1,167 @@
+import torch
+import functools
+from typing import Dict, List, Tuple, Union
+
+from cache_dit.caching.cache_contexts.cache_manager import (
+    BasicCacheConfig,
+    CachedContextManager,
+)
+from cache_dit.caching.cache_contexts.prune_context import (
+    PrunedContext,
+)
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class PrunedContextManager(CachedContextManager):
+    # Reuse CachedContextManager for Dynamic Block Prune
+
+    def __init__(self, name: str = None, **kwargs):
+        super().__init__(name, **kwargs)
+        # Overwrite for Dynamic Block Prune
+        self._current_context: PrunedContext = None
+        self._cached_context_manager: Dict[str, PrunedContext] = {}
+
+    # Overwrite for Dynamic Block Prune
+    def new_context(self, *args, **kwargs) -> PrunedContext:
+        if self._persistent_context:
+            cache_config: BasicCacheConfig = kwargs.get("cache_config", None)
+            assert (
+                cache_config is not None
+                and cache_config.num_inference_steps is not None
+            ), (
+                "When persistent_context is True, num_inference_steps "
+                "must be set in cache_config for proper cache refreshing."
+            )
+        _context = PrunedContext(*args, **kwargs)
+        # NOTE: Patch args and kwargs for implicit refresh.
+        _context._init_args = args  # maybe empty tuple: ()
+        _context._init_kwargs = kwargs  # maybe empty dict: {}
+        self._cached_context_manager[_context.name] = _context
+        return _context
+
+    def set_context(
+        self,
+        cached_context: PrunedContext | str,
+    ) -> PrunedContext:
+        return super().set_context(cached_context)
+
+    def get_context(self, name: str = None) -> PrunedContext:
+        return super().get_context(name)
+
+    def reset_context(
+        self,
+        cached_context: PrunedContext | str,
+        *args,
+        **kwargs,
+    ) -> PrunedContext:
+        return super().reset_context(cached_context, *args, **kwargs)
+
+    # Specially for Dynamic Block Prune
+    @torch.compiler.disable
+    def add_pruned_step(self):
+        cached_context = self.get_context()
+        assert cached_context is not None, "cached_context must be set before"
+        cached_context.add_pruned_step()
+
+    @torch.compiler.disable
+    def add_pruned_block(self, num_blocks):
+        cached_context = self.get_context()
+        assert cached_context is not None, "cached_context must be set before"
+        cached_context.add_pruned_block(num_blocks)
+
+    @torch.compiler.disable
+    def add_actual_block(self, num_blocks):
+        cached_context = self.get_context()
+        assert cached_context is not None, "cached_context must be set before"
+        cached_context.add_actual_block(num_blocks)
+
+    @torch.compiler.disable
+    def get_pruned_steps(self) -> List[int]:
+        cached_context = self.get_context()
+        assert cached_context is not None, "cached_context must be set before"
+        return cached_context.get_pruned_steps()
+
+    @torch.compiler.disable
+    def get_cfg_pruned_steps(self) -> List[int]:
+        cached_context = self.get_context()
+        assert cached_context is not None, "cached_context must be set before"
+        return cached_context.get_cfg_pruned_steps()
+
+    @torch.compiler.disable
+    def get_pruned_blocks(self) -> List[int]:
+        cached_context = self.get_context()
+        assert cached_context is not None, "cached_context must be set before"
+        return cached_context.get_pruned_blocks()
+
+    @torch.compiler.disable
+    def get_actual_blocks(self) -> List[int]:
+        cached_context = self.get_context()
+        assert cached_context is not None, "cached_context must be set before"
+        return cached_context.get_actual_blocks()
+
+    @torch.compiler.disable
+    def get_cfg_pruned_blocks(self) -> List[int]:
+        cached_context = self.get_context()
+        assert cached_context is not None, "cached_context must be set before"
+        return cached_context.get_cfg_pruned_blocks()
+
+    @torch.compiler.disable
+    def get_cfg_actual_blocks(self) -> List[int]:
+        cached_context = self.get_context()
+        assert cached_context is not None, "cached_context must be set before"
+        return cached_context.get_cfg_actual_blocks()
+
+    @torch.compiler.disable
+    @functools.lru_cache(maxsize=8)
+    def get_non_prune_blocks_ids(self, num_blocks: int) -> List[int]:
+        assert num_blocks is not None, "num_blocks must be provided"
+        assert num_blocks > 0, "num_blocks must be greater than 0"
+        # Get the non-prune block ids for current context
+        # Never prune the first `Fn` and last `Bn` blocks.
+        Fn_compute_blocks_ids = list(
+            range(
+                self.Fn_compute_blocks()
+                if self.Fn_compute_blocks() < num_blocks
+                else num_blocks
+            )
+        )
+
+        Bn_compute_blocks_ids = list(
+            range(
+                num_blocks
+                - (
+                    self.Bn_compute_blocks()
+                    if self.Bn_compute_blocks() < num_blocks
+                    else num_blocks
+                ),
+                num_blocks,
+            )
+        )
+        context = self.get_context()
+        assert context is not None, "cached_context must be set before"
+
+        non_prune_blocks_ids = list(
+            set(
+                Fn_compute_blocks_ids
+                + Bn_compute_blocks_ids
+                + context.cache_config.non_prune_block_ids
+            )
+        )
+        non_prune_blocks_ids = [
+            d for d in non_prune_blocks_ids if d < num_blocks
+        ]
+        return sorted(non_prune_blocks_ids)
+
+    @torch.compiler.disable
+    def can_prune(self, *args, **kwargs) -> bool:
+        # Directly reuse can_cache for Dynamic Block Prune
+        return self.can_cache(*args, **kwargs)
+
+    @torch.compiler.disable
+    def apply_prune(
+        self, *args, **kwargs
+    ) -> Tuple[torch.Tensor, Union[torch.Tensor, None]]:
+        # Directly reuse apply_cache for Dynamic Block Prune
+        return self.apply_cache(*args, **kwargs)