Merge branch 'main' into kylesayrs/transform_permutations

Author: kylesayrs

src/compressed_tensors/compressors/model_compressors/model_compressor.py

@@ -68,6 +68,10 @@
 from transformers.file_utils import CONFIG_NAME
 
 
+if TYPE_CHECKING:
+    from compressed_tensors.compressors import BaseQuantizationCompressor
+
+
 __all__ = ["ModelCompressor", "map_module_to_scheme"]
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
@@ -257,7 +261,9 @@ def __init__(
         self.sparsity_config = sparsity_config
         self.quantization_config = quantization_config
         self.sparsity_compressor = None
-        self.quantization_compressor = None
+        self.quantization_compressor: Optional[
+            Union[BaseQuantizationCompressor, DenseCompressor]
+        ] = None
 
         if sparsity_config is not None:
             self.sparsity_compressor = BaseCompressor.load_from_registry(

src/compressed_tensors/compressors/sparse_compressors/base.py

@@ -13,8 +13,9 @@
 # limitations under the License.
 
 import logging
-from typing import Dict, Generator, Optional, Set, Tuple
+from typing import TYPE_CHECKING, Dict, Generator, Optional, Set, Tuple
 
+import torch
 from compressed_tensors.compressors.base import BaseCompressor
 from compressed_tensors.utils import (
     get_nested_mappings_from_state_dict,
@@ -26,6 +27,10 @@
 from tqdm import tqdm
 
 
+if TYPE_CHECKING:
+    from compressed_tensors.quantization import QuantizationScheme
+
+
 __all__ = ["BaseSparseCompressor"]
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
@@ -200,3 +205,16 @@ def should_compress(name: str, expanded_targets: Optional[Set[str]] = None) -> b
         return (
             name.endswith(".weight") and name[: -(len(".weight"))] in expanded_targets
         )
+
+    def decompress_module_from_state_dict(
+        self,
+        prefix: str,
+        state_dict: Dict[str, torch.Tensor],
+        scheme: "QuantizationScheme",
+    ) -> Dict[str, torch.Tensor]:
+        """
+        This function is implemented as a workaround because of how
+        `ModelCompressor.quantization_compressor` can be set to either
+        an instance of `BaseQuantizationCompressor` or `BaseSparseCompressor`.
+        """
+        return state_dict.copy()

src/compressed_tensors/quantization/lifecycle/apply.py

@@ -331,12 +331,10 @@ def find_name_or_class_matches(
         3. matches on module names
     """
     targets = sorted(targets, key=lambda x: ("re:" in x, x))
-    if isinstance(targets, Iterable):
-        matches = _find_matches(name, targets) + _find_matches(
-            module.__class__.__name__, targets, check_contains
-        )
-        matches = [match for match in matches if match is not None]
-        return matches
+    matches = _find_matches(name, targets) + _find_matches(
+        module.__class__.__name__, targets, check_contains
+    )
+    return [match for match in matches if match is not None]
 
 
 def _find_matches(

src/compressed_tensors/transform/utils/hadamard.py

@@ -25,14 +25,11 @@
 # https://github.com/scipy/scipy/blob/v1.15.2/scipy/linalg/_special_matrices.py
 def deterministic_hadamard_matrix(size: int) -> torch.Tensor:
     """
-    Construct an Hadamard matrix.
+    Construct an n-by-n Hadamard matrix, using Sylvester's construction.
+    `n` must be a power of 2.
 
-    Constructs an n-by-n Hadamard matrix, using Sylvester's
-    construction. `n` must be a power of 2.
-
-    :param size: order of the matrix; must be a power of 2
-
-    returns a (size, size) hadamard matrix
+    :param size: order of the matrix, must be a power of 2
+    :return: hadamard matrix of size `size`
     """
     if size <= 0:
         raise ValueError("Cannot construct deterministic hadamard of size <= 0")

src/compressed_tensors/utils/offload.py

@@ -87,13 +87,18 @@ def decorator(func: Callable[[Any], Any]):
         if not _has_accelerate:
 
             if fallback == "error":
-                raise ValueError(
-                    "Please install `accelerate` in order to use this function"
-                )
 
-            @wraps(func)
-            def fallback_fn(*args, **kwargs):
-                return fallback
+                @wraps(func)
+                def fallback_fn(*args, **kwargs):
+                    raise ValueError(
+                        "Please install `accelerate` in order to use this function"
+                    )
+
+            else:
+
+                @wraps(func)
+                def fallback_fn(*args, **kwargs):
+                    return fallback
 
             return fallback_fn