Skip to content
Open
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions modelopt/torch/quantization/model_calib.py
Original file line number Diff line number Diff line change
Expand Up @@ -1165,6 +1165,12 @@ def awq(
See :class:`AWQFullCalibConfig <modelopt.torch.quantization.config.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)
Expand Down
14 changes: 14 additions & 0 deletions tests/unit/torch/quantization/test_calib.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Expand Down Expand Up @@ -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()
Expand Down