[Hotfix] Implement quantization compressor methods on dense compressor

src/compressed_tensors/compressors/model_compressors/model_compressor.py

@@ -30,8 +30,12 @@
     QUANTIZATION_METHOD_NAME,
     SPARSITY_CONFIG_NAME,
 )
-from compressed_tensors.compressors.base import BaseCompressor
-from compressed_tensors.compressors.sparse_compressors import DenseCompressor
+from compressed_tensors.compressors import (
+    BaseCompressor,
+    BaseQuantizationCompressor,
+    BaseSparseCompressor,
+    DenseCompressor,
+)
 from compressed_tensors.config import CompressionFormat, SparsityCompressionConfig
 from compressed_tensors.quantization import (
     DEFAULT_QUANTIZATION_METHOD,
@@ -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, BaseSparseCompressor]
+        ] = 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()