diff --git a/modelopt/torch/quantization/calib/bias.py b/modelopt/torch/quantization/calib/bias.py index e31b18bbd1c..f1614ab28b0 100644 --- a/modelopt/torch/quantization/calib/bias.py +++ b/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}") diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index 541cf5cefdf..52ef472ba17 100644 --- a/modelopt/torch/quantization/config.py +++ b/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 diff --git a/tests/unit/torch/quantization/test_affine_quant.py b/tests/unit/torch/quantization/test_affine_quant.py index 805a796b652..e370e46d43e 100644 --- a/tests/unit/torch/quantization/test_affine_quant.py +++ b/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") diff --git a/tests/unit/torch/quantization/test_config_validation.py b/tests/unit/torch/quantization/test_config_validation.py index cd08959a976..2314dc6c112 100644 --- a/tests/unit/torch/quantization/test_config_validation.py +++ b/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 "" 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"}