WIP

Signed-off-by: Kyle Sayers <[email protected]>

Author: kylesayrs

src/llmcompressor/modifiers/quantization/calibration.py

@@ -198,6 +198,10 @@ def calibrate_activations(module: Module, value: torch.Tensor, base_name: str):
             calculate_qparams = False
         if quantization_args.strategy == QuantizationStrategy.TENSOR_GROUP:
             calculate_gparam = True
+    
+    # (..., 1, hidden_dim)
+    # this reshaping is mostly for the benefit of group quantization
+    value = value.unsqueeze(-2)
 
     call_observer(
         module=module,

src/llmcompressor/observers/base.py

@@ -10,7 +10,7 @@
 )
 from compressed_tensors.quantization.utils import is_fp4
 from compressed_tensors.registry.registry import RegistryMixin
-from compressed_tensors.utils import safe_permute
+from compressed_tensors.quantization.utils import strict_divide
 from loguru import logger
 from torch import FloatTensor, IntTensor, Tensor
 
@@ -125,8 +125,6 @@ def get_qparams(
         :return: tuple of scale and zero point based on last observed value
         """
         if observed is not None:
-            group_size = self.quantization_args.group_size
-
             if self.quantization_args.strategy == QuantizationStrategy.TENSOR:
                 # re-calculate scale and zero point, update the stored value
                 self._scale, self._zero_point = self.calculate_qparams(observed)
@@ -135,50 +133,43 @@ def get_qparams(
                 QuantizationStrategy.TENSOR_GROUP,
                 QuantizationStrategy.GROUP,
             ):
-                rows = observed.shape[0]
-                columns = observed.shape[1]
-                num_groups = int(ceil(columns / group_size))
-                if num_groups * group_size != columns:
-                    logger.bind(log_once=True).warning(
-                        "Attempting to quantize a module weight whose columns "
-                        f"({columns}) are not divisible by group_size ({group_size}). "
-                        "This scheme is not supported by vLLM, please consider "
-                        "adjusting the group_size for modules with this number of "
-                        "columns",
-                    )
+                # should be identical implementation to first half of
+                # `_process_quantization`
 
-                self._scale = torch.empty(
-                    (rows, num_groups), dtype=observed.dtype, device=observed.device
-                )
+                # get shapes
+                assert observed.ndim >= 2
+                rows, columns = observed.shape[-2:]
+                group_size = self.quantization_args.group_size
+                num_groups = strict_divide(columns, group_size)
+
+                # FP4: cast zp type
                 if is_fp4(quantization_args=self.quantization_args):
                     zp_dtype = FP8_E4M3_DATA.dtype
                 else:
                     zp_dtype = self.quantization_args.pytorch_dtype()
 
+                # allocate qparams
+                self._scale = torch.empty(
+                    (rows, num_groups), dtype=observed.dtype, device=observed.device
+                )
                 self._zero_point = torch.empty(
                     (rows, num_groups), dtype=zp_dtype, device=observed.device
                 )
 
-                # support column-order (default) quantization as well as other orderings
-                # such as activation ordering. Below checks if g_idx has initialized
-                is_column_order = g_idx is None or -1 in g_idx
-                if is_column_order:
-                    group_sizes = torch.full((num_groups,), group_size, dtype=torch.int)
-                else:
-                    group_indices, group_sizes = torch.unique(g_idx, return_counts=True)
-                    group_sizes = group_sizes[torch.argsort(group_indices)]
-
+                # permute groups
+                if g_idx is not None:
                     perm = torch.argsort(g_idx)
-                    observed = safe_permute(observed, perm, dim=1)
+                    observed = observed.index_select(-1, perm)
 
                 # TODO: experiment with vectorizing for loop for performance
+                # all reduce all dims except the last one
                 end = 0
-                for group_index, group_count in enumerate(group_sizes):
+                for group_index in range(num_groups):
                     start = end
-                    end = start + group_count
+                    end = start + group_size
                     scale, zero_point = self.get_qparams_along_dim(
-                        observed[:, start:end],
-                        0,
+                        observed[..., start:end],
+                        dim=tuple(range(observed.ndim - 1)),
                         tensor_id=group_index,
                         global_scale=global_scale,
                     )
@@ -187,21 +178,23 @@ def get_qparams(
                     self._zero_point[:, group_index] = zero_point.squeeze(1)
 
             elif self.quantization_args.strategy == QuantizationStrategy.CHANNEL:
-                # assume observed is transposed, because its the output, hence use dim 0
-                self._scale, self._zero_point = self.get_qparams_along_dim(observed, 0)
+                # all reduce all dims except the last one
+                self._scale, self._zero_point = self.get_qparams_along_dim(
+                    observed,
+                    dim=tuple(range(observed.ndim - 1)),
+                )
 
             elif self.quantization_args.strategy == QuantizationStrategy.TOKEN:
-                # use dim 1, assume the obsersed.shape = [batch, token, hidden]
-                # should be batch, token
+                # all reduce all dims except the last one
                 self._scale, self._zero_point = self.get_qparams_along_dim(
                     observed,
-                    dim={0, 1},
+                    dim=tuple(range(observed.ndim - 1)),
                 )
 
             elif self.quantization_args.strategy == QuantizationStrategy.BLOCK:
                 # Block-wise quantization: one scale/zero_point per block of shape
                 # [block_rows, block_cols]
-                rows, cols = observed.shape[:2]
+                rows, cols = observed.shape[-2:]
                 bs = self.quantization_args.block_structure
                 if not (
                     isinstance(bs, (list, tuple))