minor

Signed-off-by: realAsma <akuriparambi@nvidia.com>

Author: realAsma

modelopt/torch/quantization/calib/mse.py

@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""MSE-based calibrators for quantization."""
+"""Calibrator that returns the MSE amax of all collected tensors."""
 
 import math
 from collections.abc import Callable
@@ -28,7 +28,7 @@
 
 
 class MseCalibrator(_Calibrator):
-    """MSE amax search that minimizes error between x and quantized x."""
+    """Per-tensor and per-channel MSE amax search that minimizes error between x and quantized x."""
 
     def __init__(
         self,
@@ -40,7 +40,20 @@ def __init__(
         quant_func: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None,
         error_func: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None,
     ):
-        """Initialize MSE calibrator with initial amax and search parameters."""
+        """Initialize MSE calibrator.
+
+        Args:
+            amax: Initial amax value (required).
+            axis: Quantization axis. None means per-tensor quantization.
+            step_size: Step size for amax search. The number of steps is computed as
+                        ceil((stop_multiplier - start_multiplier) / step_size) + 1.
+            start_multiplier: Starting multiplier for amax search.
+            stop_multiplier: Ending multiplier for amax search.
+            quant_func: Function that quantizes input tensor given an amax value.
+                        Should have signature: quant_func(x, amax) -> quantized_x.
+            error_func: Function to compute error between x and xq.
+                        Default is F.mse_loss(x, xq, reduction='none').
+        """
         super().__init__(num_bits=None, axis=axis, unsigned=None)
         self._initial_amax = amax
         self._step_size = step_size
@@ -68,7 +81,11 @@ def _compute_candidate_amax(self, candidates: torch.Tensor) -> torch.Tensor:
 
     @torch.no_grad()
     def collect(self, x: torch.Tensor):
-        """Collect tensor statistics for MSE calibration."""
+        """Collect input tensor statistics and compute losses for MSE calibration.
+
+        Args:
+            x: Input tensor.
+        """
         if self._quant_func is None:
             raise RuntimeError("Quantization function not set.")
 
@@ -101,13 +118,16 @@ def collect(self, x: torch.Tensor):
                 self._losses_sum[step] += loss
 
     def reset(self):
-        """Reset collected statistics."""
+        """Reset the stored losses and amax value."""
         self._losses_sum = None
         self._candidates = None
         self._amax = None
 
     def clear(self):
-        """Clear all state including initial amax."""
+        """Clear all cached data to free GPU memory.
+
+        Call this after compute_amax() and load_calib_amax() are done.
+        """
         self._losses_sum = None
         self._candidates = None
         if self._initial_amax is not None:
@@ -116,7 +136,11 @@ def clear(self):
 
     @torch.no_grad()
     def compute_amax(self, verbose: bool = False):
-        """Compute optimal amax from collected statistics."""
+        """Return the amax value that minimizes quantization error.
+
+        Args:
+            verbose: If True, print the ratio of best_amax to initial_amax.
+        """
         if self._losses_sum is None or not any(loss is not None for loss in self._losses_sum):
             return None
 

tests/gpu/torch/quantization/test_nvfp4_static_quantizer_cuda.py

@@ -212,7 +212,29 @@ def quant_func(x, amax):
         assert amax[2] > amax[1]
         assert amax[3] > amax[2]
 
-        # Test that fp8 sweep generates quantized scales
+    def test_fp8_sweep_generates_quantized_scales(self, device):
+        """Test that the fp8 sweep produces scales that are already FP8-quantized."""
+        num_blocks = 8
+        block_size = 16
+
+        x = torch.randn(num_blocks, block_size, device=device)
+        per_block_amax = x.abs().amax(dim=-1)
+        global_amax = per_block_amax.max()
+
+        def quant_func(x, amax):
+            return static_blockwise_fp4_fake_quant(x, amax, global_amax)
+
+        cal = NVFP4MSECalibrator(
+            amax=per_block_amax,
+            global_amax=global_amax,
+            quant_func=quant_func,
+        )
+
+        cal.collect(x)
+        amax = cal.compute_amax()
+
+        # The calibrator sweeps over FP8 candidates, so the resulting scales
+        # should already be representable in FP8 (i.e., quantize-dequantize is a no-op).
         scale = amax.float() / 6.0
         scale_fp8_quant_amax = global_amax.float() / 6.0
         scale_qdq = scaled_e4m3_impl(scale, scale_fp8_quant_amax)