MSE Observer Enhancement

examples/quantization_w4a16_fp4/nvfp4/llama3_nvfp4_expanded_mse_example.py

@@ -0,0 +1,41 @@
+"""NVFP4 quantization with expanded MSE observer.
+
+The ``nvfp4_expanded_mse`` observer searches over a range of per-group
+scale expansions (1.8x down to 0.8x of the observed range) to find the
+scale that minimizes quantization error.  This typically gives better
+quality than the default minmax observer at the cost of a longer
+calibration pass.
+
+Usage::
+
+    python llama3_nvfp4_expanded_mse_example.py
+"""
+
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from llmcompressor import oneshot
+from llmcompressor.modifiers.quantization import QuantizationModifier
+
+MODEL_ID = "meta-llama/Llama-3.1-8B-Instruct"
+
+model = AutoModelForCausalLM.from_pretrained(MODEL_ID)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+
+recipe = QuantizationModifier(
+    scheme="NVFP4",
+    weight_observer="nvfp4_expanded_mse",
+    ignore=["lm_head"],
+)
+
+oneshot(
+    model=model,
+    recipe=recipe,
+    dataset="ultrachat_200k",
+    splits="train_sft",
+    num_calibration_samples=512,
+    max_seq_length=2048,
+)
+
+SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-NVFP4-expanded-mse"
+model.save_pretrained(SAVE_DIR, save_compressed=True)
+tokenizer.save_pretrained(SAVE_DIR)

src/llmcompressor/observers/imatrix.py

@@ -49,6 +49,7 @@ def __init__(self, *args, **kwargs):
         self.grid = kw.get("grid", 20)
         self.norm = kw.get("norm", 3.0)
         self.strict = kw.get("strict", False)
+        self.expand = kw.get("expand", 1.0)
 
         self._imatrix_sum: Optional[torch.Tensor] = None
         self._imatrix_count: torch.Tensor = torch.tensor(0, dtype=torch.int64)
@@ -131,6 +132,7 @@ def update_statistics_from_observed(self, observed: torch.Tensor) -> None:
             self.patience,
             self.grid,
             self.norm,
+            expand=self.expand,
             importance_weights=importance_weights,
         )
 
@@ -257,6 +259,7 @@ def _grid_search(
     patience: int,
     grid: int,
     norm: float,
+    expand: float = 1.0,
     importance_weights: Optional[torch.Tensor] = None,
 ) -> MinMaxTuple:
     """Grid search for min/max minimizing (importance-weighted) quant error.
@@ -265,8 +268,14 @@ def _grid_search(
     using FP32 scales. After optimization, global_scale is computed from the final
     min/max values in get_qparams().
     """
-    min_val = torch.amin(observed, dim=(0, -1))
-    max_val = torch.amax(observed, dim=(0, -1))
+    if (
+        args.strategy == QuantizationStrategy.TENSOR_GROUP
+        and args.scale_dtype is not None
+    ):
+        args = args.model_copy(update={"scale_dtype": None})
+
+    min_val = torch.amin(observed, dim=(0, -1)) * expand
+    max_val = torch.amax(observed, dim=(0, -1)) * expand
     best_error = torch.full(
         min_val.shape,
         torch.finfo(torch.float32).max,

src/llmcompressor/observers/mse.py

@@ -26,8 +26,11 @@ def __init__(self, *args, **kwargs):
         self.grid = observer_kwargs.get("grid", 100.0)
         self.norm = observer_kwargs.get("norm", 2.4)
         self.chunk_size = observer_kwargs.get("chunk_size", 5)
+        self.expand = observer_kwargs.get("expand", 1.0)
         if self.chunk_size <= 0:
             raise ValueError(f"chunk_size must be positive, got {self.chunk_size}")
+        if self.expand < 1.0:
+            raise ValueError(f"expand value must be at least 1.0, got {self.expand}")
 
         # Pre-create token_args to avoid patch_attr context manager
         # which causes torch.compile graph breaks
@@ -45,6 +48,7 @@ def update_statistics_from_observed(self, observed: torch.Tensor) -> None:
             self.grid,
             self.norm,
             self.chunk_size,
+            self.expand,
         )
 
 
@@ -69,6 +73,7 @@ def __init__(self, *args, **kwargs):
         self.grid = observer_kwargs.get("grid", 100.0)
         self.norm = observer_kwargs.get("norm", 2.4)
         self.chunk_size = observer_kwargs.get("chunk_size", 5)
+        self.expand = observer_kwargs.get("expand", 1.0)
         if self.chunk_size <= 0:
             raise ValueError(f"chunk_size must be positive, got {self.chunk_size}")
 
@@ -88,6 +93,7 @@ def update_statistics_from_observed(self, observed: torch.Tensor) -> None:
             self.grid,
             self.norm,
             self.chunk_size,
+            self.expand,
         )
 
         if hasattr(self, "min_vals") and self.avg_constant != 1.0:
@@ -96,3 +102,30 @@ def update_statistics_from_observed(self, observed: torch.Tensor) -> None:
 
         self.min_vals = min_vals
         self.max_vals = max_vals
+
+
+@Observer.register("nvfp4_expanded_mse")
+class NVFP4ExpandedMSEObserver(MemorylessMSEObserver):
+    """
+    MSE observer with defaults tuned for NVFP4 range expansion.
+
+    Searches from ``expand`` times the observed range down to
+    ``(1 - maxshrink) * expand`` times the observed range.
+    With the defaults below, this covers 1.8x down to ~0.8x of
+    the original per-group range in 112 search steps.
+
+    Usage::
+
+        QuantizationArgs(
+            ...
+            observer="nvfp4_expanded_mse",
+        )
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        observer_kwargs = self.args.observer_kwargs
+        self.expand = observer_kwargs.get("expand", 1.8)
+        self.maxshrink = observer_kwargs.get("maxshrink", 1 - 0.8 / 1.8)
+        self.grid = observer_kwargs.get("grid", 200.0)
+        self.patience = observer_kwargs.get("patience", 1000)

src/llmcompressor/observers/mse_quant.py

@@ -3,7 +3,7 @@
 import torch
 import torch._dynamo.config
 import torch._dynamo.decorators
-from compressed_tensors.quantization import QuantizationArgs
+from compressed_tensors.quantization import QuantizationArgs, QuantizationStrategy
 from compressed_tensors.quantization.lifecycle import fake_quantize
 from compressed_tensors.quantization.utils import calculate_qparams
 
@@ -25,6 +25,7 @@ def _grid_search_mse(
     grid: float,
     norm: float,
     chunk_size: int,
+    expand: float = 1.0,
 ) -> MinMaxTuple:
     """Find per-channel min/max ranges that minimize quantization error.
 
@@ -44,9 +45,19 @@ def _grid_search_mse(
         in shrink factors
     :param norm: exponent used when computing the error. norm = 2 approximates MSE
     :param chunk_size: number of grid steps per compiled call
+    :param expand: factor to scale the initial min/max range before searching.
+        Values > 1.0 let the search explore ranges wider than the observed
+        values (e.g. expand=2.0 starts at 2x the observed range).
     """
-    min_val = torch.amin(observed, dim=(0, -1))
-    max_val = torch.amax(observed, dim=(0, -1))
+    if (
+        args.strategy == QuantizationStrategy.TENSOR_GROUP
+        and args.scale_dtype is not None
+    ):
+        args = args.model_copy(update={"scale_dtype": None})
+        token_args = token_args.model_copy(update={"scale_dtype": None})
+
+    min_val = torch.amin(observed, dim=(0, -1)) * expand
+    max_val = torch.amax(observed, dim=(0, -1)) * expand
     best_error = torch.full_like(min_val, torch.finfo(min_val.dtype).max)
     best_min_val = min_val.clone()
     best_max_val = max_val.clone()