add get_missing_module_keys to support transformers lower bound (#479)

Author: dsikka

src/compressed_tensors/compressors/model_compressors/model_compressor.py

@@ -344,6 +344,55 @@ def __init__(
                     format, config=quantization_config
                 )
 
+    def get_missing_module_keys(self, model: Module) -> List[str]:
+        """
+        Identifies the expected missing weight keys in the compressed state_dict.
+
+        When a model undergoes sparsity or quantization compression, certain
+        weight tensors may be absent from the checkpoint by virtue of compression.
+        This function determines which weight keys are missing based on the
+        applied compression techniques.
+
+        :param model: The PyTorch model to check for missing keys.
+        :return: A list of missing keys expected in the compressed state_dict.
+        """
+        missing_keys = set()
+
+        # Determine missing keys due to 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,
+            )
+
+            missing_keys.update(
+                merge_names(target_name, "weight")
+                for target_name, _module in sparse_targets
+            )
+
+        # Determine missing keys due to pack quantization
+        if (
+            self.quantization_compressor
+            and self.quantization_config.format
+            == CompressionFormat.pack_quantized.value
+        ):
+            for scheme in self.quantization_config.config_groups.values():
+                quant_targets = match_named_modules(
+                    model=model,
+                    targets=scheme.targets,
+                    ignore=self.quantization_config.ignore,
+                )
+                missing_keys.update(
+                    merge_names(target_name, "weight")
+                    for target_name, _module in quant_targets
+                )
+
+        return list(missing_keys)
+
     def get_unexpected_file_keys(self, model: Module) -> List[str]:
         """
         Identifies extra keys introduced by the compression process in the

tests/test_compressors/model_compressors/test_model_compressor.py

@@ -450,6 +450,26 @@ 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"),
+            {"layer1.weight"},
+        )
+    ],
+)
+def test_get_missing_keys(model, sparsity_config, quantization_config, expected):
+    model_compressor = ModelCompressor(
+        sparsity_config=sparsity_config, quantization_config=quantization_config
+    )
+
+    actual = model_compressor.get_missing_module_keys(model)
+    assert len(actual) == len(expected) and all(key in actual for key in expected)
+
+
 @pytest.mark.parametrize(
     "model, sparsity_config, quantization_config, expected",
     [