address reviews

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

Author: shanjiaz

src/compressed_tensors/compressors/quantized_compressors/base.py

@@ -85,6 +85,7 @@ def compress(
         """
         uncompressed_names = list(model_state.keys())
         compressed_dict = {}
+        compressed_prefixes = set()
 
         # compress values
         desc = "Compressing with quantization"
@@ -119,11 +120,26 @@ def compress(
                     device=compression_device,
                 )
 
+                compressed_prefixes.add(prefix)
+
                 # update state dict
                 for key, value in compressed_values.items():
                     compressed_dict[prefix + key] = value.to(compression_device)
 
             else:
+                # Skip qparams already added by compress_weight
+                is_duplicate = any(
+                    name.endswith(s) and name.removesuffix(s) in compressed_prefixes
+                    for s in [
+                        "weight_scale",
+                        "weight_zero_point",
+                        "weight_global_scale",
+                        "weight_g_idx",
+                    ]
+                )
+                if is_duplicate:
+                    continue
+
                 # omit saving zero points for symmetric quantization
                 if name.endswith("weight_zero_point"):
                     module_path = name.rsplit(".", 1)[0]

src/compressed_tensors/compressors/quantized_compressors/fp4_quantized.py

@@ -21,9 +21,8 @@
     BaseQuantizationCompressor,
 )
 from compressed_tensors.config import CompressionFormat
-from compressed_tensors.quantization import QuantizationArgs, QuantizationStrategy
+from compressed_tensors.quantization import QuantizationArgs
 from compressed_tensors.quantization.lifecycle.forward import dequantize, quantize
-from compressed_tensors.quantization.utils import calculate_qparam_shape
 from torch import Tensor
 
 
@@ -73,32 +72,13 @@ def compression_param_info(
         :param quantization_args: quantization parameters for the weight
         :return: dictionary mapping compressed parameter names to shape and dtype
         """
-        output = {
+        return {
             "weight_packed": (
                 torch.Size((weight_shape[0], weight_shape[1] // 2)),
                 torch.uint8,
             ),
         }
 
-        # Add weight_scale and weight_global_scale for NVFP4/MXFP4
-        if quantization_args is not None and quantization_args.strategy in [
-            QuantizationStrategy.GROUP.value,
-            QuantizationStrategy.TENSOR_GROUP.value,
-        ]:
-            # Use centralized calculation for consistency and correctness
-            num_groups, scale_shape = calculate_qparam_shape(
-                weight_shape, quantization_args
-            )
-            output["weight_scale"] = (scale_shape, quantization_args.scale_dtype)
-
-            if quantization_args.strategy == QuantizationStrategy.TENSOR_GROUP.value:
-                output["weight_global_scale"] = (
-                    torch.Size((1,)),
-                    torch.float32,
-                )
-
-        return output
-
     def compress_scale(
         self,
         scale: Tensor,

src/compressed_tensors/quantization/lifecycle/initialize.py

@@ -35,7 +35,7 @@
 from compressed_tensors.quantization.lifecycle.forward import (
     wrap_module_forward_quantized,
 )
-from compressed_tensors.quantization.utils import strategy_cdiv
+from compressed_tensors.quantization.utils import calculate_qparam_shape, strategy_cdiv
 from compressed_tensors.utils import (
     disable_hf_hook,
     get_execution_device,
@@ -215,9 +215,10 @@ def initialize_qparams(
         if len(observed_shape) < 1:
             raise ValueError("Group quant requires at least 1 observed dimension")
 
-        group_size = quantization_args.group_size
-        num_groups = strategy_cdiv(observed_shape[-1], group_size, strategy)
-        expected_shape = (*observed_shape[:-1], num_groups)
+        # Use shared calculation to avoid floor division bugs
+        _, expected_shape = calculate_qparam_shape(
+            torch.Size(observed_shape), quantization_args
+        )
 
         # initialize activation ordering if applicable
         if actorder == ActivationOrdering.GROUP: