diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 661f429b2f5..7a01c3f4559 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -1161,10 +1161,21 @@ def awq( model: Model to be calibrated. forward_loop: A callable which takes the model as argument and forwards calibration data through the model. + algorithm: The AWQ variant to run; one of ``"awq_lite"``, + ``"awq_clip"``, or ``"awq_full"``. + + Raises: + ValueError: If ``algorithm`` is not one of the valid AWQ algorithms. See :class:`AWQFullCalibConfig ` for details on the remaining arguments. """ + if algorithm not in ["awq_lite", "awq_clip", "awq_full"]: + raise ValueError( + f"Invalid AWQ algorithm {algorithm!r}. " + "Valid algorithms are 'awq_lite', 'awq_clip', and 'awq_full'." + ) + with SequentialQuantizer.convert_to_single_quantizer(model): if algorithm in ["awq_full", "awq_lite"]: awq_lite(model, forward_loop, **kwargs) diff --git a/tests/unit/torch/quantization/test_calib.py b/tests/unit/torch/quantization/test_calib.py index 9e58185e383..1aaaabe7605 100644 --- a/tests/unit/torch/quantization/test_calib.py +++ b/tests/unit/torch/quantization/test_calib.py @@ -26,6 +26,7 @@ from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.model_calib import ( apply_pre_quant_scale_and_smooth, + awq, disable_pre_quant_scale_and_resmooth, layerwise_calibrate, ) @@ -391,6 +392,19 @@ def _forward_loop(m): ) +def test_awq_unknown_algorithm_raises(): + """awq() must fail fast on an unknown algorithm instead of silently skipping calibration.""" + model = _SimpleMLP() + + with pytest.raises(ValueError, match="Invalid AWQ algorithm 'awq-lite'"): + awq(model, algorithm="awq-lite") + + # Valid algorithm strings still go through: awq_lite without a forward_loop + # warns and skips instead of raising. + with pytest.warns(UserWarning, match="forward_loop must be provided"): + awq(model, forward_loop=None, algorithm="awq_lite") + + def test_smoothquant_enable_disable(): torch.manual_seed(1234) model = _SimpleMLP()