add fp4 test

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

Author: kylesayrs

src/llmcompressor/observers/base.py

@@ -9,6 +9,7 @@
 )
 from compressed_tensors.quantization.utils import calculate_qparams, generate_gparam
 from compressed_tensors.registry.registry import RegistryMixin
+from compressed_tensors.utils import patch_attr
 
 from llmcompressor.observers.helpers import flatten_for_calibration
 
@@ -82,9 +83,11 @@ def get_global_scale(self, observed: torch.Tensor) -> torch.nn.Parameter:
         :param observed: value being observed
         :return: calibrated global parameter
         """
-        observed = observed.reshape((1, 1, -1))  # per tensor reshape
-        min_vals, max_vals = self.get_min_max(observed)
-        return generate_gparam(min_vals, max_vals)
+        # avoid updating running min/max for global scales
+        with patch_attr(self, "min_vals", None), patch_attr(self, "max_vals", None):
+            observed = observed.reshape((1, 1, -1))  # per tensor reshape
+            min_vals, max_vals = self.get_min_max(observed)
+            return generate_gparam(min_vals, max_vals)
 
     def _get_module_param(self, name: str) -> Optional[torch.nn.Parameter]:
         if self.module is None:

tests/llmcompressor/observers/test_mse.py

@@ -73,7 +73,8 @@ def test_mse_observer_symmetric_scale_range():
 
 
 def test_mse_fp4():
-    tensor = torch.arange(24, dtype=torch.bfloat16).reshape((4, 6)) / 24
+    module = torch.nn.Linear(6, 4)
+    module.weight.data = torch.arange(24, dtype=torch.bfloat16).reshape((4, 6)) / 24
 
     weights = QuantizationArgs(
         num_bits=4,
@@ -84,8 +85,15 @@ def test_mse_fp4():
     )
 
     observer = weights.observer
-    observer = Observer.load_from_registry(observer, base_name="weight", args=weights)
-    scale, zero_point = observer(tensor)
+    observer = Observer.load_from_registry(
+        observer, base_name="weight", args=weights, module=module
+    )
 
-    qdq_tensor = fake_quantize(tensor, scale, zero_point, weights)
-    assert torch.nn.functional.mse_loss(qdq_tensor, tensor) <= 0.002
+    global_scale = observer.get_global_scale(module.weight)
+    module.weight_global_scale = global_scale
+    scale, zero_point = observer(module.weight)
+
+    qdq_tensor = fake_quantize(
+        module.weight, scale, zero_point, weights, global_scale=global_scale
+    )
+    assert torch.nn.functional.mse_loss(qdq_tensor, module.weight) <= 0.002