From 90beee8929f857fca167dbc40f4f33aae2e46942 Mon Sep 17 00:00:00 2001 From: Rahul Tuli Date: Thu, 9 Jan 2025 15:07:59 +0000 Subject: [PATCH 1/3] Turn off sparse decompression when sparse compressor is dense --- .../compressors/model_compressors/model_compressor.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/src/compressed_tensors/compressors/model_compressors/model_compressor.py b/src/compressed_tensors/compressors/model_compressors/model_compressor.py index 5ea26705..4e24b2fb 100644 --- a/src/compressed_tensors/compressors/model_compressors/model_compressor.py +++ b/src/compressed_tensors/compressors/model_compressors/model_compressor.py @@ -269,9 +269,9 @@ def compress( compressed_state_dict = state_dict - quantized_modules_to_args: Dict[ - str, QuantizationArgs - ] = map_modules_to_quant_args(model) + quantized_modules_to_args: Dict[str, QuantizationArgs] = ( + map_modules_to_quant_args(model) + ) if self.quantization_compressor is not None: compressed_state_dict = self.quantization_compressor.compress( @@ -310,7 +310,10 @@ def decompress(self, model_path: str, model: Module): model_path = get_safetensors_folder(model_path) sparse_decompressed = False - if self.sparsity_compressor is not None: + if ( + self.sparsity_compressor is not None + and self.sparsity_config.format != "dense" + ): # Sparse decompression is applied on the model_path dense_gen = self.sparsity_compressor.decompress(model_path) self._replace_weights(dense_gen, model) From 26b23f75aecc49718d50db16bfa1c2ea21fb53d1 Mon Sep 17 00:00:00 2001 From: Rahul Tuli Date: Fri, 10 Jan 2025 14:05:14 +0000 Subject: [PATCH 2/3] Update: Condiition to use enum instead of raw string --- .../compressors/model_compressors/model_compressor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/compressed_tensors/compressors/model_compressors/model_compressor.py b/src/compressed_tensors/compressors/model_compressors/model_compressor.py index 4e24b2fb..78ab6422 100644 --- a/src/compressed_tensors/compressors/model_compressors/model_compressor.py +++ b/src/compressed_tensors/compressors/model_compressors/model_compressor.py @@ -312,7 +312,7 @@ def decompress(self, model_path: str, model: Module): if ( self.sparsity_compressor is not None - and self.sparsity_config.format != "dense" + and self.sparsity_config.format != CompressionFormat.dense.value ): # Sparse decompression is applied on the model_path dense_gen = self.sparsity_compressor.decompress(model_path) From cc4f78e722d3b3ea265fc69a2d6588c37353c967 Mon Sep 17 00:00:00 2001 From: Rahul Tuli Date: Fri, 10 Jan 2025 14:09:45 +0000 Subject: [PATCH 3/3] Remove unnecesarry style change --- .../compressors/model_compressors/model_compressor.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/compressed_tensors/compressors/model_compressors/model_compressor.py b/src/compressed_tensors/compressors/model_compressors/model_compressor.py index 78ab6422..951eef1f 100644 --- a/src/compressed_tensors/compressors/model_compressors/model_compressor.py +++ b/src/compressed_tensors/compressors/model_compressors/model_compressor.py @@ -269,9 +269,9 @@ def compress( compressed_state_dict = state_dict - quantized_modules_to_args: Dict[str, QuantizationArgs] = ( - map_modules_to_quant_args(model) - ) + quantized_modules_to_args: Dict[ + str, QuantizationArgs + ] = map_modules_to_quant_args(model) if self.quantization_compressor is not None: compressed_state_dict = self.quantization_compressor.compress(