fix mse, add tests

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

Author: kylesayrs

src/llmcompressor/observers/min_max.py

@@ -44,13 +44,13 @@ def get_min_max(self, observed: torch.Tensor) -> Tuple[torch.Tensor, torch.Tenso
         if self.min_vals is not None and self.averaging_constant != 1.0:
             # FUTURE: consider scaling by num observations (first dim)
             #         rather than reducing by first dim
-            min_vals = self._lerp(min_vals, self.min_vals, self.averaging_constant)
-            max_vals = self._lerp(max_vals, self.max_vals, self.averaging_constant)
+            min_vals = self._lerp(self.min_vals, min_vals, self.averaging_constant)
+            max_vals = self._lerp(self.max_vals, max_vals, self.averaging_constant)
 
         return min_vals, max_vals
 
     def _lerp(
         self, input: torch.Tensor, end: torch.Tensor, weight: float
     ) -> torch.Tensor:
         """torch lerp_kernel is not implemeneted for all data types"""
-        return (input * weight) + (end * (1.0 - weight))
+        return (input * (1.0 - weight)) + (end * weight)

src/llmcompressor/observers/mse.py

@@ -1,8 +1,12 @@
 from typing import Optional, Tuple
 
 import torch
-from compressed_tensors.quantization.quant_args import QuantizationArgs
+from compressed_tensors.quantization.quant_args import (
+    QuantizationArgs,
+    QuantizationStrategy,
+)
 from compressed_tensors.quantization.utils import calculate_qparams
+from compressed_tensors.utils import patch_attr
 
 from llmcompressor.observers.base import Observer
 
@@ -42,6 +46,24 @@ def get_min_max(self, observed: torch.Tensor) -> Tuple[torch.Tensor, torch.Tenso
             (num_observations, *qparam_shape, group_size)
         :return: minimum value and maximum value whose shapes are (*qparam_shape, )
         """
+        min_vals, max_vals = self._mse_min_max(observed)
+
+        if self.min_vals is not None and self.averaging_constant != 1.0:
+            # FUTURE: consider scaling by num observations (first dim)
+            #         rather than reducing by first dim
+            min_vals = self._lerp(self.min_vals, min_vals, self.averaging_constant)
+            max_vals = self._lerp(self.max_vals, max_vals, self.averaging_constant)
+
+        return min_vals, max_vals
+
+    def _mse_min_max(self, observed: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Grid search for MSE-optimal min and max values
+
+        :param observed: value being observed whose shape is
+            (num_observations, *qparam_shape, group_size)
+        :return: minimum and maximum values which minimize reconstruction error
+        """
         from compressed_tensors.quantization.lifecycle import fake_quantize
 
         absolute_min_val = torch.amin(observed, dim=(0, -1))
@@ -67,13 +89,17 @@ def get_min_max(self, observed: torch.Tensor) -> Tuple[torch.Tensor, torch.Tenso
                 quantization_args=self.args,
                 global_scale=global_scale,
             )
-            q = fake_quantize(
-                observed,
-                candidate_scales,
-                candidate_zero_points,
-                self.args,
-                global_scale=global_scale,
-            )
+
+            # Note that observed.shape = (num_observations, *qparams_shape, group_size).
+            # For the purposes of fake quantization, this is equivalent to token quant
+            with patch_attr(self.args, "strategy", QuantizationStrategy.TOKEN):
+                q = fake_quantize(
+                    observed,
+                    candidate_scales.unsqueeze(-1),
+                    candidate_zero_points.unsqueeze(-1),
+                    self.args,
+                    global_scale=global_scale,
+                )
 
             q -= observed
             q.abs_()
@@ -92,3 +118,9 @@ def get_min_max(self, observed: torch.Tensor) -> Tuple[torch.Tensor, torch.Tenso
                     break
 
         return min_val, max_val
+
+    def _lerp(
+        self, input: torch.Tensor, end: torch.Tensor, weight: float
+    ) -> torch.Tensor:
+        """torch lerp_kernel is not implemeneted for all data types"""
+        return (input * (1.0 - weight)) + (end * weight)

tests/llmcompressor/observers/test_mse.py

@@ -15,30 +15,45 @@
 
 import pytest
 import torch
+from compressed_tensors.quantization import fake_quantize
 from compressed_tensors.quantization.quant_args import QuantizationArgs
 
 from llmcompressor.observers import MovingAverageMSEObserver, Observer
 
 
 @pytest.mark.parametrize(
-    "symmetric,expected_scale,expected_zero_point",
+    "strategy,symmetric,exp_loss",
     [
-        (True, 0.0078, 0),
-        (False, 0.0039, -128),
+        ("tensor", True, 4.8103e-06),
+        ("tensor", False, 1.1258e-06),
+        ("channel", True, 2.5675e-06),
+        ("channel", False, 2.3696e-07),
+        ("group", True, 3.1282e-06),
+        ("group", False, 1.3794e-07),
+        ("block", True, 2.8968e-06),
+        ("block", False, 5.6068e-07),
     ],
 )
-def test_mse_observer(symmetric, expected_scale, expected_zero_point):
-    tensor = torch.tensor([1.0, 1.0, 1.0, 1.0, 1.0])
+def test_mse_observer(strategy, symmetric, exp_loss):
+    tensor = torch.arange(24).reshape((6, 4)) / 24
     num_bits = 8
-    weights = QuantizationArgs(num_bits=num_bits, symmetric=symmetric, observer="mse")
+    weights = QuantizationArgs(
+        num_bits=num_bits,
+        strategy=strategy,
+        symmetric=symmetric,
+        group_size=(2 if strategy == "group" else None),
+        block_structure=([3, 2] if strategy == "block" else None),
+        observer="mse",
+    )
 
     observer = weights.observer
     observer = Observer.load_from_registry(observer, base_name="weight", args=weights)
-    scale, zero_point = observer(tensor)
-
     assert isinstance(observer, MovingAverageMSEObserver)
-    assert round(scale.item(), 4) == expected_scale
-    assert round(zero_point.item(), 4) == expected_zero_point
+
+    scale, zero_point = observer(tensor)
+    q_tensor = fake_quantize(tensor, scale, zero_point, weights)
+    mse_loss = torch.sum((tensor - q_tensor).abs_().pow_(2)) / tensor.numel()
+    assert mse_loss == pytest.approx(exp_loss, abs=1e-10)
 
 
 def test_mse_observer_symmetric_scale_range():