address reviewed issues

Signed-off-by: shanjiaz <zsjwpianpian@gmail.com>

Author: shanjiaz

src/compressed_tensors/compressors/model_compressors/model_compressor.py

@@ -407,7 +407,7 @@ def compress_model(self, model: Module, is_meta: bool = False):
                         state_dict,
                         names_to_scheme=module_to_scheme,
                         show_progress=False,
-                        save_device=exec_device,
+                        compression_device=exec_device,
                     )
 
                 # sparsity second
@@ -416,7 +416,6 @@ def compress_model(self, model: Module, is_meta: bool = False):
                         state_dict,
                         compression_targets=sparse_compression_targets,
                         show_progress=False,
-                        module=module,
                     )
 
                 # remove any existing parameters
@@ -700,7 +699,6 @@ def _replace_sparsity_weights(self, dense_weight_generator, model: Module):
             device = "cpu" if has_offloaded_params(module) else params_device
             delattr(module, param_name)
             requires_grad = data.dtype in (torch.float16, torch.float32, torch.bfloat16)
-
             param = torch.nn.Parameter(data.to(device), requires_grad=requires_grad)
             register_offload_parameter(module, param_name, param)
 
@@ -718,6 +716,7 @@ def _replace_weights(self, dense_weight_generator, model: Module):
             'data' is the updated param data
         :param model: The model whose weights are to be updated.
         """
+
         for mod_path, data in tqdm(dense_weight_generator, desc="Decompressing model"):
             module = operator.attrgetter(mod_path)(model)
 

src/compressed_tensors/compressors/quantized_compressors/base.py

@@ -72,7 +72,7 @@ def compress(
         model_state: Dict[str, Tensor],
         names_to_scheme: Dict[str, QuantizationScheme],
         show_progress: bool = False,
-        save_device: str = "cpu",
+        compression_device: str = "cpu",
         **kwargs,
     ) -> Dict[str, Tensor]:
         """
@@ -104,7 +104,7 @@ def compress(
 
                 # is scale does not exist, then weight cannot be compressed
                 if scale is None:
-                    compressed_dict[name] = value.to(save_device)
+                    compressed_dict[name] = value.to(compression_device)
                     continue
 
                 # compress values on meta if loading from meta otherwise on cpu (memory movement too expensive)
@@ -117,12 +117,12 @@ def compress(
                     global_scale=global_scale,
                     g_idx=g_idx,
                     quantization_args=quant_args,
-                    device=save_device,
+                    device=compression_device,
                 )
 
                 # update state dict
                 for key, value in compressed_values.items():
-                    compressed_dict[prefix + key] = value.to(save_device)
+                    compressed_dict[prefix + key] = value.to(compression_device)
 
             else:
                 # omit saving zero points for symmetric or packed quantization
@@ -133,7 +133,7 @@ def compress(
                 # TODO: does this case actually occur?
                 elif name.endswith("g_idx") and torch.any(value <= -1):
                     continue
-                compressed_dict[name] = value.to(save_device)
+                compressed_dict[name] = value.to(compression_device)
 
         return compressed_dict
 

src/compressed_tensors/compressors/quantized_compressors/pack_quantized.py

@@ -241,9 +241,7 @@ def pack_to_int32(
 
     # Use int32 here
     reshaped = value.view(rows, num_groups, pack_factor).to(torch.int32)
-
     bit_shifts = torch.arange(pack_factor, device=device, dtype=torch.int32) * num_bits
-
     packed = (reshaped << bit_shifts).sum(dim=2, dtype=torch.int32)
 
     if packed_dim == 0:

src/compressed_tensors/compressors/sparse_compressors/base.py

@@ -70,7 +70,6 @@ def compress(
         model_state: Dict[str, Tensor],
         compression_targets: Optional[Set[str]] = None,
         show_progress: bool = False,
-        module: Optional[Module] = None,
     ) -> Dict[str, Tensor]:
         """
         Compresses a dense state dict using bitmask compression
@@ -96,9 +95,7 @@ def compress(
             if prefix.endswith(".weight"):
                 prefix = prefix[: -(len(".weight"))]
 
-            compression_data = self.compress_weight(
-                prefix, value, module=module
-            )
+            compression_data = self.compress_weight(prefix, value)
 
             for key in compression_data.keys():
                 if key in compressed_dict:

src/compressed_tensors/compressors/sparse_compressors/sparse_24_bitmask.py

@@ -52,17 +52,10 @@ def compression_param_names(self) -> Tuple[str]:
             "bitmask",
         )
 
-    def compress_weight(self, name, value, *, module=None):
+    def compress_weight(self, name, value):
         bitmask_tensor = Sparse24BitMaskTensor.from_dense(
             value, self.config.sparsity_structure
         )
-        if value.device.type == "meta":
-            if module is None:
-                raise ValueError("compress_weight requires module argument when is_meta=True")
-            # Create empty parameter matching compressed shape
-            empty_weight = torch.empty_like(bitmask_tensor.compressed, device="meta")
-            module.weight = torch.nn.Parameter(empty_weight, requires_grad=False)
-
         # Normal flow: return compression dict
         return bitmask_tensor.dict(
             name_prefix=name,