[Observer] Optimize mse observer (#1450)

SUMMARY:
The calculate_mse_min_max function previously performed a full grid
search across a 0.8 × 100 = 80-point space. After discussing with Alex
and Eldar last week, we reduced max_shrink to 0.2 to improve performance
without sacrificing accuracy.

Additionally, implemented an early stopping mechanism. The function now
tracks the best quantization error seen so far and stops if no
improvement is observed over 5 consecutive steps (patience = 5).

maxshrink variable is now configurable in recipe file, and patience(for
early stop) can be passed in as well.

TEST PLAN:
All lm_eval tests were run. No regressions in accuracy were observed.
Performance improved significantly after maxshrink is updated.
**There's a 3-7 mins slow down per test switching from MinMax to MSE
observer.**

USAGE:
Tested the recipe by adding:
```yaml
observer: "mse"
observer_kwargs:
    maxshrink: 0.3
```

More details can be found in this [notion
page](https://www.notion.so/Accuracy-test-1d930c7e73f3803bb057fd17d6d45302?pvs=4)
Raw timing data are stored
[here](https://drive.google.com/drive/folders/1I69QNGKxLJJZ06k9jSw0f0BRhPchV_nt?usp=drive_link)

---------

Signed-off-by: shanjiaz <[email protected]>
Co-authored-by: Brian Dellabetta <[email protected]>
Co-authored-by: Dipika Sikka <[email protected]>

Author: 3 people

src/llmcompressor/observers/mse.py

@@ -22,17 +22,19 @@ def __init__(
         quantization_args: QuantizationArgs,
         averaging_constant: float = 0.01,
         grid: float = 100.0,
-        maxshrink: float = 0.80,
         norm: float = 2.4,
         global_scale: Optional[torch.Tensor] = None,
     ):
         super().__init__(quantization_args=quantization_args, global_scale=global_scale)
 
+        kwargs = quantization_args.observer_kwargs or {}
+        self.maxshrink = kwargs.get("maxshrink", 0.20)
+        self.patience = kwargs.get("patience", 5)
+
         self.min_val = {}
         self.max_val = {}
         self.averaging_constant = averaging_constant
         self.grid = grid
-        self.maxshrink = maxshrink
         self.norm = norm
 
     def calculate_mse_min_max(
@@ -62,6 +64,10 @@ def calculate_mse_min_max(
         )
         min_val = torch.ones_like(absolute_min_val)
         max_val = torch.zeros_like(absolute_max_val)
+
+        # Early stopping params
+        no_improve_count = 0
+
         for i in range(int(self.maxshrink * self.grid)):
             p = 1 - i / self.grid
             shrinked_min_val = p * absolute_min_val
@@ -94,6 +100,12 @@ def calculate_mse_min_max(
                 best[tmp] = err[tmp]
                 min_val[tmp] = shrinked_min_val[tmp]
                 max_val[tmp] = shrinked_max_val[tmp]
+                no_improve_count = 0
+            else:
+                no_improve_count += 1
+                if no_improve_count >= self.patience:
+                    break
+
         return min_val, max_val
 
     def calculate_qparams(

tests/e2e/vLLM/recipes/actorder/recipe_w4a16_actorder_group.yaml

@@ -11,4 +11,4 @@ quant_stage:
             strategy: "group"
             group_size: 128
             actorder: "group"
-          targets: ["Linear"]
+          targets: ["Linear"]