fix(quantization): correct the affine-quant bias config contract

modelopt/torch/quantization/calib/bias.py

@@ -23,16 +23,17 @@
 
 
 def compute_maxmin(
-    inputs: torch.Tensor, axis: int | tuple[int, ...] | None
+    inputs: torch.Tensor, axis: tuple[int, ...] | list[int] | None
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Compute the max and min values of input tensor.
 
     Args:
         inputs: Input tensor
-        axis: Axis or tuple of axes to keep. Other dims are reduced.
+        axis: Tuple of axes to reduce over. The listed dims are reduced to size 1 (kept, so the
+            result broadcasts against ``inputs``); every other dim is preserved.
             None: reduce all dimensions (per-tensor)
-            (-1,): reduce all except last dim (per-channel)
-            (-1,-3): reduce all except last and third-to-last dims (per-head per-channel)
+            (-1,): reduce the last dim
+            (-1,-3): reduce the last and third-to-last dims
 
     Returns:
         Tuple of (max_values, min_values)
@@ -52,13 +53,17 @@ def compute_maxmin(
     return max_, min_
 
 
-def compute_maxmin_bias(inputs: torch.Tensor, axis: int | tuple[int, ...] | None) -> torch.Tensor:
+def compute_maxmin_bias(
+    inputs: torch.Tensor, axis: tuple[int, ...] | list[int] | None
+) -> torch.Tensor:
     """Compute the max_min mean bias of input tensor."""
     max_, min_ = compute_maxmin(inputs, axis)
     return (max_ + min_) / 2
 
 
-def compute_mean_bias(inputs: torch.Tensor, axis: int | tuple[int, ...] | None) -> torch.Tensor:
+def compute_mean_bias(
+    inputs: torch.Tensor, axis: tuple[int, ...] | list[int] | None
+) -> torch.Tensor:
     """Compute the mean bias of input tensor."""
     if axis is None:
         reduce_axis = None
@@ -75,13 +80,14 @@ def compute_mean_bias(inputs: torch.Tensor, axis: int | tuple[int, ...] | None)
 
 
 def compute_bias(
-    inputs: torch.Tensor, axis: int | tuple[int, ...] | None, method: str = "mean"
+    inputs: torch.Tensor, axis: tuple[int, ...] | list[int] | None, method: str = "mean"
 ) -> torch.Tensor:
     """Compute the bias of input tensor. Supports mean and max_min methods."""
     if method == "mean":
         return compute_mean_bias(inputs, axis)
-    else:
+    if method == "max_min":
         return compute_maxmin_bias(inputs, axis)
+    raise ValueError(f"Unsupported bias method: {method!r}, expected 'mean' or 'max_min'")
 
 
 def subtract_bias(inputs: torch.Tensor, bias: torch.Tensor) -> torch.Tensor:
@@ -100,7 +106,7 @@ def add_bias(inputs: torch.Tensor, bias: torch.Tensor) -> torch.Tensor:
 class BiasCalibrator(_Calibrator):
     """Bias calibrator, tracks the bias of all tensors collected."""
 
-    def __init__(self, method: str = "mean", axis: int | tuple[int, ...] | None = None):
+    def __init__(self, method: str = "mean", axis: tuple[int, ...] | list[int] | None = None):
         """Initialize."""
         super().__init__(axis=axis)
         self._calib_bias = None
@@ -111,17 +117,20 @@ def __init__(self, method: str = "mean", axis: int | tuple[int, ...] | None = No
 
     def collect(self, x: torch.Tensor):
         """Compute bias of input tensor along axis."""
+        # self._axis lists the dims to REDUCE; they come back as size 1 so the bias broadcasts
+        # against the input. Every dim not listed is preserved, i.e. gets its own bias.
+        #
         # For a 4D tensor with shape [batch, heads, seq_len, hidden_dim]:
-        #   - None: reduce all dimensions (per-tensor bias)
-        #   - (-1,) or (3,): keep last dimension only (per-channel bias)
-        #   - (-1, -3) or (1, 3): keep last and third-to-last dimensions (per-head per-channel bias)
-        #   This computes separate bias per attention head and channel, which is recommended
+        #   - None: reduce all dimensions (a single per-tensor bias)
+        #   - (-1,) or (3,): reduce hidden_dim, giving a separate bias per token
+        #   - (-1, -3) or (1, 3): reduce hidden_dim and heads, giving a bias per token shared
+        #     across heads, which is recommended
         #
         # Examples:
         #   tensor.shape = (8, 12, 512, 64)  # [batch, heads, seq_len, hidden]
         #   axis=None      -> single bias value for entire tensor
-        #   axis=(-1,)     -> bias shape: (1, 1, 1, 64)
-        #   axis=(-1, -3)  -> bias shape: (1, 12, 1, 64)
+        #   axis=(-1,)     -> bias shape: (8, 12, 512, 1)
+        #   axis=(-1, -3)  -> bias shape: (8, 1, 512, 1)
 
         if self._method == "mean":
             bias_ = compute_bias(x, self._axis, self._method)
@@ -162,10 +171,10 @@ def compute_bias(self):
     def compute_dynamic_bias(self, inputs):
         """Compute dynamic bias based on current inputs."""
         if self._method == "mean":
-            # mean = (max + min) / 2
+            # bias = average(all tokens)
             return compute_bias(inputs, self._axis, method="mean")
         elif self._method == "max_min":
-            # mean = average(all tokens)
+            # bias = (max + min) / 2
             return compute_bias(inputs, self._axis, method="max_min")
         else:
             raise ValueError(f"Unknown bias method: {self._method}")

modelopt/torch/quantization/config.py

@@ -524,16 +524,24 @@ def validate_learn_amax(cls, v):
         ModeloptField(
             default=None,
             title="Bias configuration.",
-            description="""Configuration for bias handling in affine quantization. The keys are:
-            - "enable": Boolean to enable/disable bias handling, default is False
+            description="""Configuration for bias handling in affine quantization.
+
+            Bias handling is enabled by setting this field; there is no separate ``"enable"`` key.
+
+            The **int keys are the axes** to reduce over when computing the bias, and map to
+            ``None``. At least one is required. Each listed dim is reduced to size 1, so a
+            separate bias is computed for every combination of the dims *not* listed. For a
+            ``[batch, heads, seq_len, hidden]`` tensor, ``{-1: None}`` gives a bias per token and
+            ``{-1: None, -3: None}`` gives a bias per token shared across heads.
+
+            Two optional string keys are also accepted:
             - "type": Specify the type of bias ["static", "dynamic"], default is "static"
             - "method": Specify the method of bias calibration ["mean", "max_min"], default is "mean"
-            - "axis": Tuple of integers specifying axes for bias computation, default is None
 
             Examples:
-            bias = {"enable": True}
-            bias = {"enable": True, "type": "static", "axis": -1}
-            bias = {"enable": True, "type": "dynamic", "axis": (-1, -3)}
+            bias = {-1: None}
+            bias = {-1: None, "type": "static"}
+            bias = {-1: None, -3: None, "type": "dynamic", "method": "max_min"}
         """,
         )
     )
@@ -583,7 +591,10 @@ def validate_bias(cls, v):
         assert len(axis) > 0, "The axis for bias computation is not specified."
         for x in axis:
             if not isinstance(x, int):
-                raise ValueError(f"Invalid axis type {type(axis)}, expected int")
+                raise ValueError(
+                    f"Unsupported bias key {x!r}. The keys are the int axes to reduce over,"
+                    ' plus the optional "type" and "method" keys.'
+                )
 
         return v
 

tests/unit/torch/quantization/test_affine_quant.py

@@ -18,6 +18,7 @@
 import pytest
 import torch
 
+from modelopt.torch.quantization.calib.bias import compute_bias
 from modelopt.torch.quantization.config import QuantizerAttributeConfig
 from modelopt.torch.quantization.model_calib import enable_stats_collection, finish_stats_collection
 from modelopt.torch.quantization.nn.modules.tensor_quantizer import TensorQuantizer
@@ -81,3 +82,12 @@ def test_bias_static(self, num_bits, type, method, axis):
             )
         elif type == "static":
             assert kv_quantizer.bias.shape == expected_bias_shape
+
+
+def test_compute_bias_rejects_unknown_method():
+    """``compute_bias`` must fail fast like its ``collect``/``compute_dynamic_bias`` siblings.
+
+    It previously fell through to ``max_min`` for any string other than ``"mean"``.
+    """
+    with pytest.raises(ValueError, match="Unsupported bias method"):
+        compute_bias(torch.randn(4), None, method="median")

tests/unit/torch/quantization/test_config_validation.py

@@ -716,3 +716,32 @@ def test_nvfp4_four_over_six_cfg_validates(self):
     def test_nvfp4_four_over_six_cfg_needs_calibration(self):
         """The 4/6 preset is statically calibrated, so it requires calibration."""
         assert need_calibration(mtq.NVFP4_FOUR_OVER_SIX_CFG)
+
+
+class TestBiasValidation:
+    """The ``bias`` field's schema: int keys are the reduction axes, plus ``type``/``method``."""
+
+    @pytest.mark.parametrize(
+        "bias_cfg",
+        [
+            {"enable": True, "type": "static", "axis": -1},
+            {"enable": True, "type": "dynamic", "axis": (-1, -3)},
+            {"axis": -1},
+        ],
+    )
+    def test_unknown_bias_string_key_is_named_in_the_error(self, bias_cfg):
+        """An unsupported string key is reported as such, not mislabelled as a bad axis type."""
+        with pytest.raises(ValidationError) as exc_info:
+            QuantizerAttributeConfig(num_bits=8, bias=bias_cfg)
+
+        message = str(exc_info.value)
+        unknown_key = next(k for k in bias_cfg if k not in ("type", "method"))
+        assert "Unsupported bias key" in message
+        assert repr(unknown_key) in message
+        # The old message reported the type of the whole key list instead of the offending key.
+        assert "<class 'list'>" not in message
+
+    def test_bias_axis_keys_are_accepted(self):
+        """Int keys are the reduction axes and must keep validating."""
+        cfg = QuantizerAttributeConfig(bias={-1: None, -3: None, "type": "static"})
+        assert cfg.bias == {-1: None, -3: None, "type": "static"}