From 342e3ca2adbd12807d521e3f22af661e2aeba7c3 Mon Sep 17 00:00:00 2001 From: Kyle Sayers Date: Tue, 10 Jun 2025 16:26:48 -0400 Subject: [PATCH 1/3] safeguard against removing parameters not immediately on the module Signed-off-by: Kyle Sayers --- .../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 c5c10845..42e7b30a 100644 --- a/src/compressed_tensors/compressors/model_compressors/model_compressor.py +++ b/src/compressed_tensors/compressors/model_compressors/model_compressor.py @@ -483,7 +483,7 @@ def decompress_model(self, model: Module): # remove any existing parameters exec_device = get_execution_device(module) offload_device = get_offloaded_device(module) - for name, _ in list(module.named_parameters()): + for name, _ in list(module.named_parameters(recurse=False)): delete_offload_parameter(module, name) # replace with decompressed parameters From d54d00b4d85dbbfdb5312f26fe2b8462e4d612ac Mon Sep 17 00:00:00 2001 From: Kyle Sayers Date: Wed, 9 Jul 2025 22:07:25 -0400 Subject: [PATCH 2/3] add safeguard to compress method Signed-off-by: Kyle Sayers --- .../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 42e7b30a..727156ed 100644 --- a/src/compressed_tensors/compressors/model_compressors/model_compressor.py +++ b/src/compressed_tensors/compressors/model_compressors/model_compressor.py @@ -421,7 +421,7 @@ def compress_model(self, model: Module): # remove any existing parameters offload_device = get_offloaded_device(module) - for name, _ in list(module.named_parameters()): + for name, _ in list(module.named_parameters(recurse=False)): delete_offload_parameter(module, name) # replace with compressed parameters From fad52c915d2205f1d561fe0e63aea253c7cf9c37 Mon Sep 17 00:00:00 2001 From: Kyle Sayers Date: Wed, 9 Jul 2025 22:12:37 -0400 Subject: [PATCH 3/3] safeguard against state dict creation Signed-off-by: Kyle Sayers --- .../compressors/model_compressors/model_compressor.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/src/compressed_tensors/compressors/model_compressors/model_compressor.py b/src/compressed_tensors/compressors/model_compressors/model_compressor.py index 727156ed..3637ecd3 100644 --- a/src/compressed_tensors/compressors/model_compressors/model_compressor.py +++ b/src/compressed_tensors/compressors/model_compressors/model_compressor.py @@ -400,7 +400,10 @@ def compress_model(self, model: Module): # in the future, support compression on same device with align_module_device(module, execution_device=exec_device): - state_dict = module.state_dict(prefix=f"{prefix}.") + state_dict = { + f"{prefix}.{name}": param + for name, param in module.named_parameters(recurse=False) + } # quantization first if prefix in module_to_scheme: @@ -458,7 +461,10 @@ def decompress_model(self, model: Module): if prefix in module_to_scheme or prefix in sparse_compression_targets: # in the future, support decompression on same device with align_module_device(module, execution_device="cpu"): - state_dict = module.state_dict(prefix=f"{prefix}.") + state_dict = { + f"{prefix}.{name}": param + for name, param in module.named_parameters(recurse=False) + } # sparsity first if prefix in sparse_compression_targets: