[NVFP4] Update to use tensor_group strategy; update observers (#1484)

SUMMARY:
- Requires neuralmagic/compressed-tensors#325
- Uses the new `tensor_group` strategy for nvfp4a16 quantization 
- Removes global_scale as an observer class parameter and passes in as a
function call, similar to g_idx

Author: dsikka

src/llmcompressor/modifiers/quantization/calibration.py

@@ -7,7 +7,7 @@
     QuantizationStatus,
 )
 from compressed_tensors.quantization.lifecycle.forward import forward_quantize
-from compressed_tensors.quantization.utils import is_fp4, is_kv_cache_quant_scheme
+from compressed_tensors.quantization.utils import is_kv_cache_quant_scheme
 from compressed_tensors.utils import align_module_device, update_parameter_data
 from loguru import logger
 from torch.nn import Module
@@ -54,14 +54,9 @@ def initialize_observer(
     quantization_args = getattr(quantization_scheme, arg_name, None)
     # dont need observers for dynamic
     if quantization_args is not None and not quantization_args.dynamic:
-        global_scale = getattr(module, f"{base_name}_global_scale", None)
-        if global_scale is not None:
-            assert base_name == "weight" and is_fp4(quantization_args=quantization_args)
-
         observer = Observer.load_from_registry(
             quantization_args.observer,
             quantization_args=quantization_args,
-            global_scale=global_scale,
         )
         module.register_module(f"{base_name}_observer", observer)
 
@@ -80,15 +75,19 @@ def call_observer(module: Module, base_name: str, value: Optional[torch.Tensor]
         if base_name == "weight":
             value = module.weight
             g_idx = getattr(module, "weight_g_idx", None)
+            global_scale = getattr(module, f"{base_name}_global_scale", None)
         elif value is not None:
             g_idx = None
+            global_scale = None
         else:
             raise ValueError(
                 "Must provide a value to observe if not using weight observer"
             )
 
         observer = getattr(module, f"{base_name}_observer")
-        updated_scale, updated_zero_point = observer(value, g_idx=g_idx)
+        updated_scale, updated_zero_point = observer(
+            value, g_idx=g_idx, global_scale=global_scale
+        )
 
         # update scale and zero point
         update_parameter_data(module, updated_scale, f"{base_name}_scale")

src/llmcompressor/observers/base.py

@@ -27,28 +27,32 @@ class Observer(Module, RegistryMixin):
     def __init__(
         self,
         quantization_args: QuantizationArgs,
-        global_scale: Optional[torch.Tensor] = None,
     ):
         self.quantization_args: QuantizationArgs = quantization_args
         super().__init__()
-        self.global_scale: Optional[torch.Tensor] = global_scale
         self._scale = None
         self._zero_point = None
         self._num_observed_tokens = None
 
     @torch.no_grad()
     def forward(
-        self, observed: Tensor, g_idx: Optional[Tensor] = None
+        self,
+        observed: Tensor,
+        g_idx: Optional[Tensor] = None,
+        global_scale: Optional[Tensor] = None,
     ) -> Tuple[FloatTensor, IntTensor]:
         """
         maps directly to get_qparams
         :param observed: optional observed tensor from which to calculate
             quantization parameters
         :param g_idx: optional mapping from column index to group index
+        :param global_scale: optional scale to further scale local quantization scales
         :return: tuple of scale and zero point based on last observed value
         """
         self.record_observed_tokens(observed)
-        return self.get_qparams(observed=observed, g_idx=g_idx)
+        return self.get_qparams(
+            observed=observed, g_idx=g_idx, global_scale=global_scale
+        )
 
     def calculate_qparams(
         self,
@@ -73,6 +77,7 @@ def get_qparams(
         self,
         observed: Optional[Tensor] = None,
         g_idx: Optional[Tensor] = None,
+        global_scale: Optional[Tensor] = None,
     ) -> Tuple[FloatTensor, IntTensor]:
         """
         Convenience function to wrap overwritten calculate_qparams
@@ -82,6 +87,7 @@ def get_qparams(
         :param observed: optional observed tensor to calculate quantization parameters
             from
         :param g_idx: optional mapping from column index to group index
+        :param global_scale: optional scale to further scale local quantization scales
         :return: tuple of scale and zero point based on last observed value
         """
         if observed is not None:
@@ -91,7 +97,10 @@ def get_qparams(
                 # re-calculate scale and zero point, update the stored value
                 self._scale, self._zero_point = self.calculate_qparams(observed)
 
-            elif self.quantization_args.strategy == QuantizationStrategy.GROUP:
+            elif self.quantization_args.strategy in (
+                QuantizationStrategy.TENSOR_GROUP,
+                QuantizationStrategy.GROUP,
+            ):
                 rows = observed.shape[0]
                 columns = observed.shape[1]
                 num_groups = int(ceil(columns / group_size))
@@ -128,6 +137,7 @@ def get_qparams(
                         observed[:, start:end],
                         0,
                         tensor_id=group_index,
+                        global_scale=global_scale
                     )
 
                     self._scale[:, group_index] = scale.squeeze(1)
@@ -160,14 +170,18 @@ def get_qparams_along_dim(
         observed,
         dim: Union[int, Iterable[int]],
         tensor_id: Optional[Any] = None,
+        global_scale: Optional[Tensor] = None,
     ):
         if isinstance(dim, int):
             dim = [dim]
         dim = set(dim)
 
         reduce_dims = tuple(idx for idx in range(observed.ndim) if idx not in dim)
         return self.calculate_qparams(
-            observed, reduce_dims=reduce_dims, tensor_id=tensor_id
+            observed,
+            reduce_dims=reduce_dims,
+            tensor_id=tensor_id,
+            global_scale=global_scale,
         )
 
     def record_observed_tokens(self, batch_tensor: Tensor):

src/llmcompressor/observers/min_max.py

@@ -22,9 +22,8 @@ def __init__(
         self,
         quantization_args: QuantizationArgs,
         averaging_constant: float = 0.01,
-        global_scale: Optional[torch.Tensor] = None,
     ):
-        super().__init__(quantization_args=quantization_args, global_scale=global_scale)
+        super().__init__(quantization_args=quantization_args)
 
         self.min_val = {}
         self.max_val = {}
@@ -35,6 +34,7 @@ def calculate_qparams(
         observed: torch.Tensor,
         reduce_dims: Optional[Tuple[int]] = None,
         tensor_id: Optional[Any] = None,
+        global_scale: Optional[torch.Tensor] = None,
     ) -> Tuple[torch.FloatTensor, torch.IntTensor]:
         """
         Updates the observed min and max using a moving average smoothed by the
@@ -46,6 +46,7 @@ def calculate_qparams(
             reduced dimensions
         :param tensor_id: Optional id if different ranges of observed tensors are
             passed, useful for sharding tensors by group_size
+        :param global_scale: optional scale to further scale local quantization scales
         :return: tuple of scale and zero point derived from the observed tensor
         """
         tensor_id = tensor_id or "default"
@@ -62,7 +63,7 @@ def calculate_qparams(
                 min_vals=min_val,
                 max_vals=max_val,
                 quantization_args=self.quantization_args,
-                global_scale=self.global_scale,
+                global_scale=global_scale,
             )
 
         running_min_val = self.min_val.get(tensor_id, None)
@@ -86,18 +87,25 @@ def calculate_qparams(
             min_vals=updated_min_val,
             max_vals=updated_max_val,
             quantization_args=self.quantization_args,
-            global_scale=self.global_scale,
+            global_scale=global_scale,
         )
 
     def get_qparams_along_dim(
-        self, observed: torch.Tensor, dim: int, tensor_id: Optional[Any] = None
+        self,
+        observed: torch.Tensor,
+        dim: int,
+        tensor_id: Optional[Any] = None,
+        global_scale: Optional[torch.Tensor] = None,
     ):
         """
         Calculate quantization parameters along the specified dimension
         """
         reduce_dims = tuple(idx for idx in range(observed.ndim) if idx != dim)
         return self.calculate_qparams(
-            observed, reduce_dims=reduce_dims, tensor_id=tensor_id
+            observed,
+            reduce_dims=reduce_dims,
+            tensor_id=tensor_id,
+            global_scale=global_scale,
         )
 
     def reset(self):

src/llmcompressor/observers/mse.py

@@ -23,9 +23,8 @@ def __init__(
         averaging_constant: float = 0.01,
         grid: float = 100.0,
         norm: float = 2.4,
-        global_scale: Optional[torch.Tensor] = None,
     ):
-        super().__init__(quantization_args=quantization_args, global_scale=global_scale)
+        super().__init__(quantization_args=quantization_args)
 
         kwargs = quantization_args.observer_kwargs or {}
         self.maxshrink = kwargs.get("maxshrink", 0.20)
@@ -41,6 +40,7 @@ def calculate_mse_min_max(
         self,
         observed: Tensor,
         reduce_dims: Optional[Tuple[int]] = None,
+        global_scale: Optional[torch.Tensor] = None,
     ):
         """
         Computes the mse-clipped min and max values of the observed tensor by
@@ -49,6 +49,7 @@ def calculate_mse_min_max(
         :param observed: observed tensor to calculate quantization parameters for
         :param reduce_dims: optional tuple of dimensions to reduce along,
             returned values will be shaped (1,) along the reduced dimensions
+        :param global_scale: optional scale to further scale local quantization scales
         :return: tuple of min and max values derived from the observed tensor
         """
         from compressed_tensors.quantization.lifecycle import fake_quantize
@@ -77,14 +78,14 @@ def calculate_mse_min_max(
                 min_vals=shrinked_min_val,
                 max_vals=shrinked_max_val,
                 quantization_args=self.quantization_args,
-                global_scale=self.global_scale,
+                global_scale=global_scale,
             )
             q = fake_quantize(
                 observed,
                 candidate_scales,
                 candidate_zero_points,
                 self.quantization_args,
-                global_scale=self.global_scale,
+                global_scale=global_scale,
             )
 
             q -= observed
@@ -113,6 +114,7 @@ def calculate_qparams(
         observed: Tensor,
         reduce_dims: Optional[Tuple[int]] = None,
         tensor_id: Optional[Any] = None,
+        global_scale: Optional[torch.Tensor] = None,
     ) -> Tuple[FloatTensor, IntTensor]:
         """
         Updates the mse-clipped min and max values of the observed tensor using
@@ -124,6 +126,7 @@ def calculate_qparams(
             reduced dimensions
         :param tensor_id: Optional id if different ranges of observed tensors are
             passed, useful for sharding tensors by group_size
+        :param global_scale: optional scale to further scale local quantization scales
         :return: tuple of scale and zero point derived from the observed tensor
         """
         min_val, max_val = self.calculate_mse_min_max(observed, reduce_dims)
@@ -150,15 +153,22 @@ def calculate_qparams(
             min_vals=updated_min_val,
             max_vals=updated_max_val,
             quantization_args=self.quantization_args,
-            global_scale=self.global_scale,
+            global_scale=global_scale,
         )
 
     def get_qparams_along_dim(
-        self, observed, dim: int, tensor_id: Optional[Any] = None
+        self,
+        observed,
+        dim: int,
+        tensor_id: Optional[Any] = None,
+        global_scale: Optional[torch.Tensor] = None,
     ):
         reduce_dims = tuple(idx for idx in range(observed.ndim) if idx != dim)
         return self.calculate_qparams(
-            observed, reduce_dims=reduce_dims, tensor_id=tensor_id
+            observed,
+            reduce_dims=reduce_dims,
+            tensor_id=tensor_id,
+            global_scale=global_scale,
         )
 
     def reset(self):