[ModelCompressor] Remove missing keys and missing modules (#462)

* remove missing keys and missing modules

* format, remove tests

* format

Author: dsikka

src/compressed_tensors/compressors/model_compressors/model_compressor.py

@@ -50,7 +50,6 @@
     get_offloaded_device,
     get_safetensors_folder,
     has_offloaded_params,
-    merge_names,
     register_offload_parameter,
     update_parameter_data,
 )
@@ -321,112 +320,6 @@ def __init__(
                     format, config=quantization_config
                 )
 
-    # ----- used by hf quantizer ----- #
-
-    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
-        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

@@ -253,61 +253,6 @@ def forward(self, x):
         return x
 
 
-@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",
-    [
-        (
-            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)
-
-
 def _create_dummy_checkpoint(state_dict, save_dir, model_compressor):
     save_dir = Path(save_dir)
     save_dir.mkdir(parents=True, exist_ok=True)