address reviews

Signed-off-by: shanjiaz <[email protected]>

Author: shanjiaz

src/llmcompressor/modifiers/quantization/calibration.py

@@ -13,6 +13,7 @@
 from compressed_tensors.utils import (
     align_module_device,
     delete_offload_parameter,
+    register_offload_parameter,
     update_offload_parameter,
 )
 from loguru import logger
@@ -131,21 +132,23 @@ def call_observer(
             # register or update scale & zero_point parameters (supports block shapes)
             scale_name = f"{base_name}_scale"
             zp_name = f"{base_name}_zero_point"
-            for name, value in [
+            for name, param_value in [
                 (scale_name, updated_scale),
                 (zp_name, updated_zero_point),
             ]:
                 if (
                     not hasattr(module, name)
-                    or getattr(module, name).shape != value.shape
+                    or getattr(module, name).shape != param_value.shape
                 ):
                     if hasattr(module, name):
                         delete_offload_parameter(module, name)
-                    register_offload_parameter(module
-                        name, torch.nn.Parameter(value.clone(), requires_grad=False)
+                    register_offload_parameter(
+                        module,
+                        name,
+                        torch.nn.Parameter(param_value.clone(), requires_grad=False),
                     )
                 else:
-                    update_offload_parameter(module, name, value)
+                    update_offload_parameter(module, name, param_value)
 
 
 def update_weight_global_scale(module: Module):

src/llmcompressor/observers/base.py

@@ -209,13 +209,22 @@ def get_qparams(
                 block_rows, block_cols = bs
                 num_br = int(ceil(rows / block_rows))
                 num_bc = int(ceil(cols / block_cols))
+
                 # allocate per-block scale and zero_point
                 self._scale = torch.empty(
                     (num_br, num_bc), dtype=observed.dtype, device=observed.device
                 )
+                
+                # Use same dtype logic as GROUP strategy for zero_point
+                if is_fp4(quantization_args=self.quantization_args):
+                    zp_dtype = FP8_E4M3_DATA.dtype
+                else:
+                    zp_dtype = self.quantization_args.pytorch_dtype()
+                
                 self._zero_point = torch.empty(
-                    (num_br, num_bc), dtype=observed.dtype, device=observed.device
+                    (num_br, num_bc), dtype=zp_dtype, device=observed.device
                 )
+
                 # compute qparams for each block
                 for i in range(num_br):
                     r0 = i * block_rows