Replace COMPRESSION_PARAM_NAMES with Abstract Property (#249)

* Add compression_param_names abstract property

* Review Comments
- Get rid of COMPRESSION_PARAM_NAMES
- Enforce implementation of compression_param_names
- Make compression_param_names immutable by using tuple instead of list

Author: rahul-tuli

src/compressed_tensors/compressors/base.py

@@ -77,6 +77,15 @@ def compression_param_info(
         """
         raise NotImplementedError()
 
+    @property
+    @abstractmethod
+    def compression_param_names(self) -> Tuple[str]:
+        """
+        Returns a tuple of compression parameter names introduced by
+        the compressor during compression
+        """
+        raise NotImplementedError()
+
     @abstractmethod
     def compress(
         self,

src/compressed_tensors/compressors/quantized_compressors/base.py

@@ -144,7 +144,7 @@ def decompress(
 
     def _decompress_from_path(self, path_to_model, names_to_scheme, device):
         weight_mappings = get_nested_weight_mappings(
-            path_to_model, self.COMPRESSION_PARAM_NAMES
+            path_to_model, self.compression_param_names
         )
         for weight_name in weight_mappings.keys():
             weight_data = {}
@@ -161,7 +161,7 @@ def _decompress_from_path(self, path_to_model, names_to_scheme, device):
 
     def _decompress_from_state_dict(self, state_dict, names_to_scheme):
         weight_mappings = get_nested_mappings_from_state_dict(
-            state_dict, self.COMPRESSION_PARAM_NAMES
+            state_dict, self.compression_param_names
         )
         for weight_name in weight_mappings.keys():
             weight_data = {}

src/compressed_tensors/compressors/quantized_compressors/naive_quantized.py

@@ -41,12 +41,18 @@ class NaiveQuantizationCompressor(BaseQuantizationCompressor):
     type to the type specified by the layer's QuantizationArgs.
     """
 
-    COMPRESSION_PARAM_NAMES = [
-        "weight",
-        "weight_scale",
-        "weight_zero_point",
-        "weight_g_idx",
-    ]
+    @property
+    def compression_param_names(self) -> Tuple[str]:
+        """
+        Returns a tuple of compression parameter names introduced by
+        the compressor during compression
+        """
+        return (
+            "weight",
+            "weight_scale",
+            "weight_zero_point",
+            "weight_g_idx",
+        )
 
     def compression_param_info(
         self,

src/compressed_tensors/compressors/quantized_compressors/pack_quantized.py

@@ -36,13 +36,19 @@ class PackedQuantizationCompressor(BaseQuantizationCompressor):
     Compresses a quantized model by packing every eight 4-bit weights into an int32
     """
 
-    COMPRESSION_PARAM_NAMES = [
-        "weight_packed",
-        "weight_scale",
-        "weight_zero_point",
-        "weight_g_idx",
-        "weight_shape",
-    ]
+    @property
+    def compression_param_names(self) -> Tuple[str]:
+        """
+        Returns a tuple of compression parameter names introduced by
+        the compressor during compression
+        """
+        return (
+            "weight_packed",
+            "weight_scale",
+            "weight_zero_point",
+            "weight_g_idx",
+            "weight_shape",
+        )
 
     def compression_param_info(
         self,

src/compressed_tensors/compressors/sparse_compressors/base.py

@@ -30,8 +30,7 @@
 class BaseSparseCompressor(BaseCompressor):
     """
     Base class representing a sparse compression algorithm. Each child class should
-    implement compression_param_info, compress_weight and decompress_weight; child
-    classes should also define COMPRESSION_PARAM_NAMES.
+    implement compression_param_names, compress_weight and decompress_weight;
 
     Compressors support compressing/decompressing a full module state dict or a single
     quantized PyTorch leaf module.
@@ -113,7 +112,7 @@ def decompress(
         """
         weight_mappings, ignored_params = get_nested_weight_mappings(
             path_to_model_or_tensors,
-            self.COMPRESSION_PARAM_NAMES,
+            self.compression_param_names,
             return_unmatched_params=True,
         )
         for weight_name in weight_mappings.keys():

src/compressed_tensors/compressors/sparse_compressors/dense.py

@@ -25,6 +25,14 @@ class DenseCompressor(BaseCompressor):
     Identity compressor for dense models, returns the original state_dict
     """
 
+    @property
+    def compression_param_names(self) -> Tuple[str]:
+        """
+        Returns a tuple of compression parameter names introduced by
+        the compressor during compression
+        """
+        return ()
+
     def compress(self, model_state: Dict[str, Tensor], **kwargs) -> Dict[str, Tensor]:
         return model_state
 

src/compressed_tensors/compressors/sparse_compressors/sparse_24_bitmask.py

@@ -40,11 +40,17 @@ class Sparse24BitMaskCompressor(BaseSparseCompressor):
     values tensor, with their locations stored in a 2d bitmask
     """
 
-    COMPRESSION_PARAM_NAMES = [
-        "shape",
-        "compressed",
-        "bitmask",
-    ]
+    @property
+    def compression_param_names(self) -> Tuple[str]:
+        """
+        Returns a tuple of compression parameter names introduced by
+        the compressor during compression
+        """
+        return (
+            "shape",
+            "compressed",
+            "bitmask",
+        )
 
     def compress_weight(self, name, value):
         bitmask_tensor = Sparse24BitMaskTensor.from_dense(

src/compressed_tensors/compressors/sparse_compressors/sparse_bitmask.py

@@ -38,7 +38,13 @@ class BitmaskCompressor(BaseSparseCompressor):
     values tensor, with their locations stored in a 2d bitmask
     """
 
-    COMPRESSION_PARAM_NAMES = ["shape", "compressed", "bitmask", "row_offsets"]
+    @property
+    def compression_param_names(self) -> Tuple[str]:
+        """
+        Returns a tuple of compression parameter names introduced by
+        the compressor during compression
+        """
+        return ("shape", "compressed", "bitmask", "row_offsets")
 
     def compress_weight(self, name, value):
         bitmask_tensor = BitmaskTensor.from_dense(value)

src/compressed_tensors/compressors/sparse_quantized_compressors/marlin_24.py

@@ -42,8 +42,6 @@ class Marlin24Compressor(BaseCompressor):
     Marlin24 kernel. Decompression is not implemented for this compressor.
     """
 
-    COMPRESSION_PARAM_NAMES = ["weight_packed", "scale_packed", "meta"]
-
     @staticmethod
     def validate_quant_compatability(
         model_quant_args: Dict[str, QuantizationArgs]
@@ -105,6 +103,14 @@ def validate_sparsity_structure(name: str, weight: Tensor) -> bool:
 
         return True
 
+    @property
+    def compression_param_names(self) -> Tuple[str]:
+        """
+        Returns a tuple of compression parameter names introduced by
+        the compressor during compression
+        """
+        return ("weight_packed", "scale_packed", "meta")
+
     def compress(
         self,
         model_state: Dict[str, Tensor],

src/compressed_tensors/utils/safetensors_load.py

@@ -16,7 +16,7 @@
 import os
 import re
 import struct
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Dict, Iterable, Optional, Tuple, Union
 
 from safetensors import safe_open
 from torch import Tensor
@@ -180,7 +180,9 @@ def get_weight_mappings(path_to_model_or_tensors: str) -> Dict[str, str]:
 
 
 def get_nested_weight_mappings(
-    model_path: str, params_to_nest: List[str], return_unmatched_params: bool = False
+    model_path: str,
+    params_to_nest: Iterable[str],
+    return_unmatched_params: bool = False,
 ) -> Union[NestedWeightMappingType, Tuple[NestedWeightMappingType, WeightMappingType]]:
     """
     Takes a path to a state dict saved in safetensors format and returns a nested
@@ -211,7 +213,7 @@ def get_nested_weight_mappings(
 
     :param model_path: Path to the safetensors state dict, must contain either a
         single safetensors file or multiple files with an index.
-    :param params_to_nest: List of parameter names to nest.
+    :param params_to_nest: Iterable of parameter names to nest.
     :param return_unmatched_params: If True, return a second dictionary containing
         the remaining parameters that were not matched to the params_to_nest.
     :return:
@@ -247,7 +249,7 @@ def get_nested_weight_mappings(
 
 
 def get_nested_mappings_from_state_dict(
-    state_dict, params_to_nest
+    state_dict, params_to_nest: Iterable[str]
 ) -> NestedWeightMappingType:
     """
     Takes a state dict and returns a nested mapping from uncompressed
@@ -262,7 +264,7 @@ def get_nested_mappings_from_state_dict(
     }
 
     :param state_dict: state dict of the model
-    :param params_to_nest: List of parameter names to nest.
+    :param params_to_nest: Iterable of parameter names to nest.
     :return: Nested mapping of parameterized layer names to the value of
         each layer's compression parameters.
     """