add back get missing keys to support transformers lower bound (#475)

Author: dsikka

src/compressed_tensors/compressors/model_compressors/model_compressor.py

@@ -52,6 +52,7 @@
     get_offloaded_device,
     get_safetensors_folder,
     has_offloaded_params,
+    merge_names,
     patch_attr,
     register_offload_parameter,
     update_parameter_data,
@@ -343,6 +344,61 @@ def __init__(
                     format, config=quantization_config
                 )
 
+    def get_unexpected_file_keys(self, model: Module) -> List[str]:
+        """
+        Identifies extra keys introduced by the compression process in the
+        compressed state_dict that are not expected by the model graph.
+
+        During sparsity or quantization compression, additional metadata or
+        auxiliary parameters may be stored in the checkpoint, which do not
+        correspond to any parameter in the original model. These keys are
+        typically introduced to support the reconstruction of compressed weights.
+
+        For example, Sparse24Bitmask compression may introduce keys such as
+        'compressed', 'bitmask', and 'shape' in the checkpoint, which are
+        not part of the original model parameters.
+
+        :param model: The PyTorch model to check for unexpected keys.
+        :return: A list of extra keys introduced by the compression process
+                that are not expected by the model.
+        """
+
+        unexpected_keys = set()
+
+        # Identify unexpected keys from sparsity compression
+        if (
+            self.sparsity_compressor
+            and self.sparsity_config.format != CompressionFormat.dense.value
+        ):
+            sparse_targets = match_named_modules(
+                model=model,
+                targets=self.sparsity_config.targets,
+                ignore=self.sparsity_config.ignore,
+            )
+            unexpected_keys.update(
+                merge_names(target_name, param)
+                for target_name, _module in sparse_targets
+                for param in self.sparsity_compressor.compression_param_names
+            )
+
+        # Identify unexpected keys from quantization compression
+        if self.quantization_compressor:
+            for scheme in self.quantization_config.config_groups.values():
+                quant_targets = match_named_modules(
+                    model=model,
+                    targets=scheme.targets,
+                    ignore=self.quantization_config.ignore,
+                )
+                for quant_compressor in self.quantization_compressor.values():
+                    unexpected_keys.update(
+                        merge_names(target_name, param)
+                        for target_name, _module in quant_targets
+                        for param in quant_compressor.compression_param_names
+                        if param != "weight"
+                    )
+
+        return list(unexpected_keys)
+
     # ----- model memory compression/decompression pathways ----- #
 
     def compress_model(self, model: Module):

tests/test_compressors/model_compressors/test_model_compressor.py

@@ -450,6 +450,41 @@ def test_multiple_quant_compressors():
     assert all(format in compressor.quantization_compressor for format in formats)
 
 
+@pytest.mark.parametrize(
+    "model, sparsity_config, quantization_config, expected",
+    [
+        (
+            TwoLayerModel(),
+            get_bitmask_sparsity_config(targets=["re:.*layer1$"]),
+            create_quantization_config(bits=8, type="int", strategy="channel"),
+            {
+                f"{layer}.{suffix}"
+                for layer, suffixes in {
+                    "layer1": [
+                        "shape",
+                        "row_offsets",
+                        "weight_zero_point",
+                        "weight_g_idx",
+                        "bitmask",
+                        "weight_scale",
+                        "compressed",
+                    ],
+                    "layer2": ["weight_scale", "weight_zero_point", "weight_g_idx"],
+                }.items()
+                for suffix in suffixes
+            },
+        )
+    ],
+)
+def test_get_unexpected_keys(model, sparsity_config, quantization_config, expected):
+    model_compressor = ModelCompressor(
+        sparsity_config=sparsity_config, quantization_config=quantization_config
+    )
+
+    actual = model_compressor.get_unexpected_file_keys(model)
+    assert len(actual) == len(expected) and all(key in actual for key in expected)
+
+
 @pytest.mark.parametrize(
     "model_stub,comp_stub",
     [