cache-dit 0.2.22__py3-none-any.whl → 0.2.24__py3-none-any.whl

cache_dit/__init__.py

@@ -9,8 +9,8 @@ from cache_dit.cache_factory import enable_cache
 from cache_dit.cache_factory import cache_type
 from cache_dit.cache_factory import block_range
 from cache_dit.cache_factory import CacheType
+from cache_dit.cache_factory import BlockAdapter
 from cache_dit.cache_factory import ForwardPattern
-from cache_dit.cache_factory import BlockAdapterParams
 from cache_dit.compile import set_compile_configs
 from cache_dit.utils import summary
 from cache_dit.utils import strify
@@ -19,8 +19,6 @@ from cache_dit.logger import init_logger
 NONE = CacheType.NONE
 DBCache = CacheType.DBCache
 
-BlockAdapter = BlockAdapterParams
-
 Forward_Pattern_0 = ForwardPattern.Pattern_0
 Forward_Pattern_1 = ForwardPattern.Pattern_1
 Forward_Pattern_2 = ForwardPattern.Pattern_2

cache_dit/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.2.22'
-__version_tuple__ = version_tuple = (0, 2, 22)
+__version__ = version = '0.2.24'
+__version_tuple__ = version_tuple = (0, 2, 24)
 
 __commit_id__ = commit_id = None

cache_dit/cache_factory/__init__.py

@@ -2,7 +2,7 @@ from cache_dit.cache_factory.forward_pattern import ForwardPattern
 from cache_dit.cache_factory.cache_types import CacheType
 from cache_dit.cache_factory.cache_types import cache_type
 from cache_dit.cache_factory.cache_types import block_range
-from cache_dit.cache_factory.cache_adapters import BlockAdapterParams
+from cache_dit.cache_factory.cache_adapters import BlockAdapter
 from cache_dit.cache_factory.cache_adapters import UnifiedCacheAdapter
 from cache_dit.cache_factory.cache_interface import enable_cache
 from cache_dit.cache_factory.utils import load_options

cache_dit/cache_factory/cache_adapters.py

@@ -5,7 +5,7 @@ import unittest
 import functools
 import dataclasses
 
-from typing import Any
+from typing import Any, Tuple, List, Optional
 from contextlib import ExitStack
 from diffusers import DiffusionPipeline
 from cache_dit.cache_factory.patch.flux import (
@@ -23,28 +23,251 @@ logger = init_logger(__name__)
 
 
 @dataclasses.dataclass
-class BlockAdapterParams:
+class BlockAdapter:
     pipe: DiffusionPipeline = None
     transformer: torch.nn.Module = None
     blocks: torch.nn.ModuleList = None
     # transformer_blocks, blocks, etc.
     blocks_name: str = None
     dummy_blocks_names: list[str] = dataclasses.field(default_factory=list)
+    # flags to control auto block adapter
+    auto: bool = False
+    allow_prefixes: List[str] = dataclasses.field(
+        default_factory=lambda: [
+            "transformer",
+            "single_transformer",
+            "blocks",
+            "layers",
+        ]
+    )
+    check_prefixes: bool = True
+    allow_suffixes: List[str] = dataclasses.field(
+        default_factory=lambda: ["TransformerBlock"]
+    )
+    check_suffixes: bool = False
+    blocks_policy: str = dataclasses.field(
+        default="max", metadata={"allowed_values": ["max", "min"]}
+    )
+
+    def __post_init__(self):
+        self.maybe_apply_patch()
+
+    def maybe_apply_patch(self):
+        # Process some specificial cases, specific for transformers
+        # that has different forward patterns between single_transformer_blocks
+        # and transformer_blocks , such as Flux (diffusers < 0.35.0).
+        if self.transformer.__class__.__name__.startswith("Flux"):
+            self.transformer = maybe_patch_flux_transformer(
+                self.transformer,
+                blocks=self.blocks,
+            )
+
+    @staticmethod
+    def auto_block_adapter(
+        adapter: "BlockAdapter",
+        forward_pattern: Optional[ForwardPattern] = None,
+    ) -> "BlockAdapter":
+        assert adapter.auto, (
+            "Please manually set `auto` to True, or, manually "
+            "set all the transformer blocks configuration."
+        )
+        assert adapter.pipe is not None, "adapter.pipe can not be None."
+        pipe = adapter.pipe
+
+        assert hasattr(pipe, "transformer"), "pipe.transformer can not be None."
+
+        transformer = pipe.transformer
+
+        # "transformer_blocks", "blocks", "single_transformer_blocks", "layers"
+        blocks, blocks_name = BlockAdapter.find_blocks(
+            transformer=transformer,
+            allow_prefixes=adapter.allow_prefixes,
+            allow_suffixes=adapter.allow_suffixes,
+            check_prefixes=adapter.check_prefixes,
+            check_suffixes=adapter.check_suffixes,
+            blocks_policy=adapter.blocks_policy,
+            forward_pattern=forward_pattern,
+        )
+
+        return BlockAdapter(
+            pipe=pipe,
+            transformer=transformer,
+            blocks=blocks,
+            blocks_name=blocks_name,
+        )
 
-    def check_adapter_params(self) -> bool:
+    @staticmethod
+    def check_block_adapter(adapter: "BlockAdapter") -> bool:
         if (
-            isinstance(self.pipe, DiffusionPipeline)
-            and self.transformer is not None
-            and self.blocks is not None
-            and isinstance(self.blocks, torch.nn.ModuleList)
+            isinstance(adapter.pipe, DiffusionPipeline)
+            and adapter.transformer is not None
+            and adapter.blocks is not None
+            and adapter.blocks_name is not None
+            and isinstance(adapter.blocks, torch.nn.ModuleList)
         ):
             return True
+
+        logger.warning("Check block adapter failed!")
         return False
 
+    @staticmethod
+    def find_blocks(
+        transformer: torch.nn.Module,
+        allow_prefixes: List[str] = [
+            "transformer",
+            "single_transformer",
+            "blocks",
+            "layers",
+        ],
+        allow_suffixes: List[str] = [
+            "TransformerBlock",
+        ],
+        check_prefixes: bool = True,
+        check_suffixes: bool = False,
+        **kwargs,
+    ) -> Tuple[torch.nn.ModuleList, str]:
+        # Check prefixes
+        if check_prefixes:
+            blocks_names = []
+            for attr_name in dir(transformer):
+                for prefix in allow_prefixes:
+                    if attr_name.startswith(prefix):
+                        blocks_names.append(attr_name)
+        else:
+            blocks_names = dir(transformer)
+
+        # Check ModuleList
+        valid_names = []
+        valid_count = []
+        forward_pattern = kwargs.get("forward_pattern", None)
+        for blocks_name in blocks_names:
+            if blocks := getattr(transformer, blocks_name, None):
+                if isinstance(blocks, torch.nn.ModuleList):
+                    block = blocks[0]
+                    block_cls_name = block.__class__.__name__
+                    # Check suffixes
+                    if isinstance(block, torch.nn.Module) and (
+                        any(
+                            (
+                                block_cls_name.endswith(allow_suffix)
+                                for allow_suffix in allow_suffixes
+                            )
+                        )
+                        or (not check_suffixes)
+                    ):
+                        # May check forward pattern
+                        if forward_pattern is not None:
+                            if BlockAdapter.match_blocks_pattern(
+                                blocks,
+                                forward_pattern,
+                                logging=False,
+                            ):
+                                valid_names.append(blocks_name)
+                                valid_count.append(len(blocks))
+                        else:
+                            valid_names.append(blocks_name)
+                            valid_count.append(len(blocks))
+
+        if not valid_names:
+            raise ValueError(
+                "Auto selected transformer blocks failed, please set it manually."
+            )
+
+        final_name = valid_names[0]
+        final_count = valid_count[0]
+        block_policy = kwargs.get("blocks_policy", "max")
+
+        for blocks_name, count in zip(valid_names, valid_count):
+            blocks = getattr(transformer, blocks_name)
+            logger.info(
+                f"Auto selected transformer blocks: {blocks_name}, "
+                f"class: {blocks[0].__class__.__name__}, "
+                f"num blocks: {count}"
+            )
+            if block_policy == "max":
+                if final_count < count:
+                    final_count = count
+                    final_name = blocks_name
+            else:
+                if final_count > count:
+                    final_count = count
+                    final_name = blocks_name
+
+        final_blocks = getattr(transformer, final_name)
+
+        logger.info(
+            f"Final selected transformer blocks: {final_name}, "
+            f"class: {final_blocks[0].__class__.__name__}, "
+            f"num blocks: {final_count}, block_policy: {block_policy}."
+        )
+
+        return final_blocks, final_name
+
+    @staticmethod
+    def match_block_pattern(
+        block: torch.nn.Module,
+        forward_pattern: ForwardPattern,
+    ) -> bool:
+        assert (
+            forward_pattern.Supported
+            and forward_pattern in ForwardPattern.supported_patterns()
+        ), f"Pattern {forward_pattern} is not support now!"
+
+        forward_parameters = set(
+            inspect.signature(block.forward).parameters.keys()
+        )
+        num_outputs = str(
+            inspect.signature(block.forward).return_annotation
+        ).count("torch.Tensor")
+
+        in_matched = True
+        out_matched = True
+        if num_outputs > 0 and len(forward_pattern.Out) != num_outputs:
+            # output pattern not match
+            out_matched = False
+
+        for required_param in forward_pattern.In:
+            if required_param not in forward_parameters:
+                in_matched = False
+
+        return in_matched and out_matched
+
+    @staticmethod
+    def match_blocks_pattern(
+        transformer_blocks: torch.nn.ModuleList,
+        forward_pattern: ForwardPattern,
+        logging: bool = True,
+    ) -> bool:
+        assert (
+            forward_pattern.Supported
+            and forward_pattern in ForwardPattern.supported_patterns()
+        ), f"Pattern {forward_pattern} is not support now!"
+
+        assert isinstance(transformer_blocks, torch.nn.ModuleList)
+
+        pattern_matched_states = []
+        for block in transformer_blocks:
+            pattern_matched_states.append(
+                BlockAdapter.match_block_pattern(
+                    block,
+                    forward_pattern,
+                )
+            )
+
+        pattern_matched = all(pattern_matched_states)  # all block match
+        if pattern_matched and logging:
+            block_cls_name = transformer_blocks[0].__class__.__name__
+            logger.info(
+                f"Match Block Forward Pattern: {block_cls_name}, {forward_pattern}"
+                f"\nIN:{forward_pattern.In}, OUT:{forward_pattern.Out})"
+            )
+
+        return pattern_matched
+
 
 @dataclasses.dataclass
 class UnifiedCacheParams:
-    adapter_params: BlockAdapterParams = None
+    block_adapter: BlockAdapter = None
     forward_pattern: ForwardPattern = ForwardPattern.Pattern_0
 
 
@@ -85,7 +308,7 @@ class UnifiedCacheAdapter:
 
             assert isinstance(pipe.transformer, FluxTransformer2DModel)
             return UnifiedCacheParams(
-                adapter_params=BlockAdapterParams(
+                block_adapter=BlockAdapter(
                     pipe=pipe,
                     transformer=pipe.transformer,
                     blocks=(
@@ -102,7 +325,7 @@ class UnifiedCacheAdapter:
 
             assert isinstance(pipe.transformer, MochiTransformer3DModel)
             return UnifiedCacheParams(
-                adapter_params=BlockAdapterParams(
+                block_adapter=BlockAdapter(
                     pipe=pipe,
                     transformer=pipe.transformer,
                     blocks=pipe.transformer.transformer_blocks,
@@ -116,7 +339,7 @@ class UnifiedCacheAdapter:
 
             assert isinstance(pipe.transformer, CogVideoXTransformer3DModel)
             return UnifiedCacheParams(
-                adapter_params=BlockAdapterParams(
+                block_adapter=BlockAdapter(
                     pipe=pipe,
                     transformer=pipe.transformer,
                     blocks=pipe.transformer.transformer_blocks,
@@ -136,7 +359,7 @@ class UnifiedCacheAdapter:
                 (WanTransformer3DModel, WanVACETransformer3DModel),
             )
             return UnifiedCacheParams(
-                adapter_params=BlockAdapterParams(
+                block_adapter=BlockAdapter(
                     pipe=pipe,
                     transformer=pipe.transformer,
                     blocks=pipe.transformer.blocks,
@@ -150,7 +373,7 @@ class UnifiedCacheAdapter:
 
             assert isinstance(pipe.transformer, HunyuanVideoTransformer3DModel)
             return UnifiedCacheParams(
-                adapter_params=BlockAdapterParams(
+                block_adapter=BlockAdapter(
                     pipe=pipe,
                     blocks=(
                         pipe.transformer.transformer_blocks
@@ -166,7 +389,7 @@ class UnifiedCacheAdapter:
 
             assert isinstance(pipe.transformer, QwenImageTransformer2DModel)
             return UnifiedCacheParams(
-                adapter_params=BlockAdapterParams(
+                block_adapter=BlockAdapter(
                     pipe=pipe,
                     transformer=pipe.transformer,
                     blocks=pipe.transformer.transformer_blocks,
@@ -180,7 +403,7 @@ class UnifiedCacheAdapter:
 
             assert isinstance(pipe.transformer, LTXVideoTransformer3DModel)
             return UnifiedCacheParams(
-                adapter_params=BlockAdapterParams(
+                block_adapter=BlockAdapter(
                     pipe=pipe,
                     transformer=pipe.transformer,
                     blocks=pipe.transformer.transformer_blocks,
@@ -194,7 +417,7 @@ class UnifiedCacheAdapter:
 
             assert isinstance(pipe.transformer, AllegroTransformer3DModel)
             return UnifiedCacheParams(
-                adapter_params=BlockAdapterParams(
+                block_adapter=BlockAdapter(
                     pipe=pipe,
                     transformer=pipe.transformer,
                     blocks=pipe.transformer.transformer_blocks,
@@ -208,7 +431,7 @@ class UnifiedCacheAdapter:
 
             assert isinstance(pipe.transformer, CogView3PlusTransformer2DModel)
             return UnifiedCacheParams(
-                adapter_params=BlockAdapterParams(
+                block_adapter=BlockAdapter(
                     pipe=pipe,
                     transformer=pipe.transformer,
                     blocks=pipe.transformer.transformer_blocks,
@@ -222,7 +445,7 @@ class UnifiedCacheAdapter:
 
             assert isinstance(pipe.transformer, CogView4Transformer2DModel)
             return UnifiedCacheParams(
-                adapter_params=BlockAdapterParams(
+                block_adapter=BlockAdapter(
                     pipe=pipe,
                     transformer=pipe.transformer,
                     blocks=pipe.transformer.transformer_blocks,
@@ -236,7 +459,7 @@ class UnifiedCacheAdapter:
 
             assert isinstance(pipe.transformer, CosmosTransformer3DModel)
             return UnifiedCacheParams(
-                adapter_params=BlockAdapterParams(
+                block_adapter=BlockAdapter(
                     pipe=pipe,
                     transformer=pipe.transformer,
                     blocks=pipe.transformer.transformer_blocks,
@@ -250,7 +473,7 @@ class UnifiedCacheAdapter:
 
             assert isinstance(pipe.transformer, EasyAnimateTransformer3DModel)
             return UnifiedCacheParams(
-                adapter_params=BlockAdapterParams(
+                block_adapter=BlockAdapter(
                     pipe=pipe,
                     transformer=pipe.transformer,
                     blocks=pipe.transformer.transformer_blocks,
@@ -264,7 +487,7 @@ class UnifiedCacheAdapter:
 
             assert isinstance(pipe.transformer, SkyReelsV2Transformer3DModel)
             return UnifiedCacheParams(
-                adapter_params=BlockAdapterParams(
+                block_adapter=BlockAdapter(
                     pipe=pipe,
                     transformer=pipe.transformer,
                     blocks=pipe.transformer.blocks,
@@ -278,7 +501,7 @@ class UnifiedCacheAdapter:
 
             assert isinstance(pipe.transformer, SD3Transformer2DModel)
             return UnifiedCacheParams(
-                adapter_params=BlockAdapterParams(
+                block_adapter=BlockAdapter(
                     pipe=pipe,
                     transformer=pipe.transformer,
                     blocks=pipe.transformer.transformer_blocks,
@@ -294,13 +517,13 @@ class UnifiedCacheAdapter:
     def apply(
         cls,
         pipe: DiffusionPipeline = None,
-        adapter_params: BlockAdapterParams = None,
+        block_adapter: BlockAdapter = None,
         forward_pattern: ForwardPattern = ForwardPattern.Pattern_0,
         **cache_context_kwargs,
     ) -> DiffusionPipeline:
         assert (
-            pipe is not None or adapter_params is not None
-        ), "pipe or adapter_params can not both None!"
+            pipe is not None or block_adapter is not None
+        ), "pipe or block_adapter can not both None!"
 
         if pipe is not None:
             if cls.is_supported(pipe):
@@ -310,7 +533,7 @@ class UnifiedCacheAdapter:
                 )
                 params = cls.get_params(pipe)
                 return cls.cachify(
-                    params.adapter_params,
+                    params.block_adapter,
                     forward_pattern=params.forward_pattern,
                     **cache_context_kwargs,
                 )
@@ -324,7 +547,7 @@ class UnifiedCacheAdapter:
                 "Adapting cache acceleration using custom BlockAdapter!"
             )
             return cls.cachify(
-                adapter_params,
+                block_adapter,
                 forward_pattern=forward_pattern,
                 **cache_context_kwargs,
             )
@@ -332,22 +555,48 @@ class UnifiedCacheAdapter:
     @classmethod
     def cachify(
         cls,
-        adapter_params: BlockAdapterParams,
+        block_adapter: BlockAdapter,
         *,
         forward_pattern: ForwardPattern = ForwardPattern.Pattern_0,
         **cache_context_kwargs,
     ) -> DiffusionPipeline:
-        if adapter_params.check_adapter_params():
-            assert isinstance(adapter_params.blocks, torch.nn.ModuleList)
+
+        if block_adapter.auto:
+            block_adapter = BlockAdapter.auto_block_adapter(
+                block_adapter,
+                forward_pattern,
+            )
+
+        if BlockAdapter.check_block_adapter(block_adapter):
             # Apply cache on pipeline: wrap cache context
-            cls.create_context(adapter_params.pipe, **cache_context_kwargs)
+            cls.create_context(
+                block_adapter.pipe,
+                **cache_context_kwargs,
+            )
             # Apply cache on transformer: mock cached transformer blocks
             cls.mock_blocks(
-                adapter_params,
+                block_adapter,
                 forward_pattern=forward_pattern,
             )
+            cls.patch_params(
+                block_adapter,
+                forward_pattern=forward_pattern,
+                **cache_context_kwargs,
+            )
+        return block_adapter.pipe
 
-        return adapter_params.pipe
+    @classmethod
+    def patch_params(
+        cls,
+        block_adapter: BlockAdapter,
+        forward_pattern: ForwardPattern = None,
+        **cache_context_kwargs,
+    ):
+        block_adapter.transformer._forward_pattern = forward_pattern
+        block_adapter.transformer._cache_context_kwargs = cache_context_kwargs
+        block_adapter.pipe.__class__._cache_context_kwargs = (
+            cache_context_kwargs
+        )
 
     @classmethod
     def has_separate_cfg(
@@ -407,28 +656,21 @@ class UnifiedCacheAdapter:
 
         pipe.__class__.__call__ = new_call
         pipe.__class__._is_cached = True
-        pipe.__class__._cache_options = cache_kwargs
         return pipe
 
     @classmethod
     def mock_blocks(
         cls,
-        adapter_params: BlockAdapterParams,
+        block_adapter: BlockAdapter,
         forward_pattern: ForwardPattern = ForwardPattern.Pattern_0,
     ) -> torch.nn.Module:
-        if getattr(adapter_params.transformer, "_is_cached", False):
-            return adapter_params.transformer
-
-        # Firstly, process some specificial cases (TODO: more patches)
-        if adapter_params.transformer.__class__.__name__.startswith("Flux"):
-            adapter_params.transformer = maybe_patch_flux_transformer(
-                adapter_params.transformer,
-                blocks=adapter_params.blocks,
-            )
+
+        if getattr(block_adapter.transformer, "_is_cached", False):
+            return block_adapter.transformer
 
         # Check block forward pattern matching
-        assert cls.match_pattern(
-            adapter_params.blocks,
+        assert BlockAdapter.match_blocks_pattern(
+            block_adapter.blocks,
             forward_pattern=forward_pattern,
         ), (
             "No block forward pattern matched, "
@@ -439,22 +681,17 @@ class UnifiedCacheAdapter:
         cached_blocks = torch.nn.ModuleList(
             [
                 DBCachedTransformerBlocks(
-                    adapter_params.blocks,
-                    transformer=adapter_params.transformer,
+                    block_adapter.blocks,
+                    transformer=block_adapter.transformer,
                     forward_pattern=forward_pattern,
                 )
             ]
         )
         dummy_blocks = torch.nn.ModuleList()
 
-        original_forward = adapter_params.transformer.forward
+        original_forward = block_adapter.transformer.forward
 
-        assert isinstance(adapter_params.dummy_blocks_names, list)
-        if adapter_params.blocks_name is None:
-            adapter_params.blocks_name = cls.find_blocks_name(
-                adapter_params.transformer
-            )
-            assert adapter_params.blocks_name is not None
+        assert isinstance(block_adapter.dummy_blocks_names, list)
 
         @functools.wraps(original_forward)
         def new_forward(self, *args, **kwargs):
@@ -462,11 +699,11 @@ class UnifiedCacheAdapter:
                 stack.enter_context(
                     unittest.mock.patch.object(
                         self,
-                        adapter_params.blocks_name,
+                        block_adapter.blocks_name,
                         cached_blocks,
                     )
                 )
-                for dummy_name in adapter_params.dummy_blocks_names:
+                for dummy_name in block_adapter.dummy_blocks_names:
                     stack.enter_context(
                         unittest.mock.patch.object(
                             self,
@@ -476,71 +713,9 @@ class UnifiedCacheAdapter:
                     )
                 return original_forward(*args, **kwargs)
 
-        adapter_params.transformer.forward = new_forward.__get__(
-            adapter_params.transformer
+        block_adapter.transformer.forward = new_forward.__get__(
+            block_adapter.transformer
         )
-        adapter_params.transformer._is_cached = True
+        block_adapter.transformer._is_cached = True
 
-        return adapter_params.transformer
-
-    @classmethod
-    def match_pattern(
-        cls,
-        transformer_blocks: torch.nn.ModuleList,
-        forward_pattern: ForwardPattern = ForwardPattern.Pattern_0,
-    ) -> bool:
-        pattern_matched_states = []
-
-        assert (
-            forward_pattern.Supported
-            and forward_pattern in ForwardPattern.supported_patterns()
-        ), f"Pattern {forward_pattern} is not support now!"
-
-        for block in transformer_blocks:
-            forward_parameters = set(
-                inspect.signature(block.forward).parameters.keys()
-            )
-            num_outputs = str(
-                inspect.signature(block.forward).return_annotation
-            ).count("torch.Tensor")
-
-            in_matched = True
-            out_matched = True
-            if num_outputs > 0 and len(forward_pattern.Out) != num_outputs:
-                # output pattern not match
-                out_matched = False
-
-            for required_param in forward_pattern.In:
-                if required_param not in forward_parameters:
-                    in_matched = False
-
-            pattern_matched_states.append(in_matched and out_matched)
-
-        pattern_matched = all(pattern_matched_states)  # all block match
-        if pattern_matched:
-            block_cls_name = transformer_blocks[0].__class__.__name__
-            logger.info(
-                f"Match Block Forward Pattern: {block_cls_name}, {forward_pattern}"
-                f"\nIN:{forward_pattern.In}, OUT:{forward_pattern.Out})"
-            )
-
-        return pattern_matched
-
-    @classmethod
-    def find_blocks_name(cls, transformer):
-        blocks_name = None
-        allow_prefixes = ["transformer", "blocks"]
-        for attr_name in dir(transformer):
-            if blocks_name is None:
-                for prefix in allow_prefixes:
-                    # transformer_blocks, blocks
-                    if attr_name.startswith(prefix):
-                        blocks_name = attr_name
-                        logger.info(f"Auto selected blocks name: {blocks_name}")
-                        # only find one transformer blocks name
-                        break
-        if blocks_name is None:
-            logger.warning(
-                "Auto selected blocks name failed, please set it manually."
-            )
-        return blocks_name
+        return block_adapter.transformer

cache_dit/cache_factory/cache_blocks.py

@@ -1,5 +1,6 @@
 import inspect
 import torch
+import torch.distributed as dist
 
 from cache_dit.cache_factory import cache_context
 from cache_dit.cache_factory import ForwardPattern
@@ -179,10 +180,15 @@ class DBCachedTransformerBlocks(torch.nn.Module):
     @torch.compiler.disable
     def _is_parallelized(self):
         # Compatible with distributed inference.
-        return all(
+        return any(
             (
-                self.transformer is not None,
-                getattr(self.transformer, "_is_parallelized", False),
+                all(
+                    (
+                        self.transformer is not None,
+                        getattr(self.transformer, "_is_parallelized", False),
+                    )
+                ),
+                (dist.is_initialized() and dist.get_world_size() > 1),
             )
         )