update to test 8 bits; update kwargs

Author: dsikka

src/compressed_tensors/compressors/base.py

@@ -45,7 +45,7 @@ def compress(self, model_state: Dict[str, Tensor], **kwargs) -> Dict[str, Tensor
         raise NotImplementedError()
 
     def decompress(
-        self, path_to_model_or_tensors: str, device: str = "cpu"
+        self, path_to_model_or_tensors: str, device: str = "cpu", **kwargs
     ) -> Generator[Tuple[str, Tensor], None, None]:
         """
         Reads a compressed state dict located at path_to_model_or_tensors

src/compressed_tensors/compressors/model_compressor.py

@@ -25,7 +25,7 @@
     SPARSITY_CONFIG_NAME,
 )
 from compressed_tensors.compressors import Compressor
-from compressed_tensors.config import CompressionFormat, SparsityCompressionConfig
+from compressed_tensors.config import SparsityCompressionConfig
 from compressed_tensors.quantization import (
     QuantizationConfig,
     QuantizationStatus,
@@ -254,15 +254,9 @@ def decompress(self, model_path: str, model: Module):
         if self.quantization_compressor is not None:
             names_to_scheme = apply_quantization_config(model, self.quantization_config)
             load_pretrained_quantization(model, model_path)
-            if (
-                self.quantization_config.format
-                == CompressionFormat.pack_quantized.value
-            ):
-                dense_gen = self.quantization_compressor.decompress(
-                    model_path, names_to_scheme=names_to_scheme
-                )
-            else:
-                dense_gen = self.quantization_compressor.decompress(model_path)
+            dense_gen = self.quantization_compressor.decompress(
+                model_path, names_to_scheme=names_to_scheme
+            )
             self._replace_weights(dense_gen, model)
 
             def update_status(module):

tests/test_compressors/test_pack_quant.py

@@ -32,10 +32,10 @@
 from safetensors.torch import save_file
 
 
-def get_dummy_quant_config():
+def get_dummy_quant_config(num_bits=4):
     config_groups = {
         "group_1": QuantizationScheme(
-            targets=["Linear"], weights=QuantizationArgs(num_bits=4)
+            targets=["Linear"], weights=QuantizationArgs(num_bits=num_bits)
         ),
     }
     ignore = ["lm_head"]
@@ -106,7 +106,8 @@ def test_repack(value):
     assert torch.equal(value, unpacked)
 
 
-def test_reload_match(tmp_path):
+@pytest.mark.parametrize("num_bits", [4, 8])
+def test_reload_match(tmp_path, num_bits):
     dense_state_dict = {
         "dummy.weight": torch.rand((511, 350)),
         "dummy.weight_scale": torch.tensor(0.01, dtype=torch.float32),
@@ -115,11 +116,12 @@ def test_reload_match(tmp_path):
         "dummy2.weight_scale": torch.tensor(0.02, dtype=torch.float32),
         "dummy2.weight_zero_point": torch.tensor(15, dtype=torch.int32),
     }
+    print("num bits", num_bits)
     names_to_scheme = {
-        "dummy": QuantizationArgs(num_bits=4),
-        "dummy2": QuantizationArgs(num_bits=4),
+        "dummy": QuantizationArgs(num_bits=num_bits),
+        "dummy2": QuantizationArgs(num_bits=num_bits),
     }
-    quant_config = get_dummy_quant_config()
+    quant_config = get_dummy_quant_config(num_bits)
 
     compressor = PackedQuantizationCompressor(config=quant_config)
     quantized_modules_to_args = {