From d98e59a43395a27ecc750da1f5e4a78b3f240646 Mon Sep 17 00:00:00 2001 From: n1ck-guo Date: Fri, 7 Aug 2026 14:03:34 +0800 Subject: [PATCH] fix performance bug Signed-off-by: n1ck-guo --- auto_round/algorithms/quantization/base.py | 13 ------------- auto_round/algorithms/quantization/rtn/quantizer.py | 3 ++- .../quantization/sign_roundv2/quantizer.py | 4 ++-- 3 files changed, 4 insertions(+), 16 deletions(-) diff --git a/auto_round/algorithms/quantization/base.py b/auto_round/algorithms/quantization/base.py index c83d805d5b..70b3299e41 100644 --- a/auto_round/algorithms/quantization/base.py +++ b/auto_round/algorithms/quantization/base.py @@ -198,19 +198,6 @@ def quantize_layer_outside_block( the valid-token loss mask. ``None`` for RTN fallback. """ self._quantize_layer_via_rtn(layer, disable_opt_rtn=disable_opt_rtn) - """Quantize a single layer outside a transformer block using RTN fallback. - Args: - layer: The layer module to quantize. Must have a - ``global_name`` attribute for model re-insertion. - fp_inputs: Optional FP calibration inputs; unused in base RTN. - q_inputs: Optional quantized activations; unused in base RTN. - disable_opt_rtn: ``True`` skips optimized-RTN scale/zp search. - ``None`` defers to ``self.config.disable_opt_rtn``. - valid_token_mask: Per-sample masks; unused in base RTN. - input_ids: Original FP calibration inputs, same as ``fp_inputs`` - when ``q_inputs`` is ``None``; unused in base RTN. - """ - self._quantize_layer_via_rtn(layer, disable_opt_rtn=disable_opt_rtn) @torch.no_grad() def _quantize_layer_via_rtn(self, layer: "torch.nn.Module", disable_opt_rtn: "bool | None" = None) -> None: diff --git a/auto_round/algorithms/quantization/rtn/quantizer.py b/auto_round/algorithms/quantization/rtn/quantizer.py index 6fee985fd2..e4c6466e2e 100644 --- a/auto_round/algorithms/quantization/rtn/quantizer.py +++ b/auto_round/algorithms/quantization/rtn/quantizer.py @@ -19,6 +19,7 @@ from auto_round.algorithms.registry import register_pipeline_member from auto_round.logger import logger from auto_round.utils import ( + SUPPORTED_LAYER_TYPES, check_to_quantized, ) @@ -99,7 +100,7 @@ def collect_imatrix(module, input, output): handles = [] for _, module in model.named_modules(): - if check_to_quantized(module): + if isinstance(module, SUPPORTED_LAYER_TYPES) and check_to_quantized(module): handles.append(module.register_forward_hook(collect_imatrix)) return handles diff --git a/auto_round/algorithms/quantization/sign_roundv2/quantizer.py b/auto_round/algorithms/quantization/sign_roundv2/quantizer.py index 4c306f5ccc..e99494fed0 100644 --- a/auto_round/algorithms/quantization/sign_roundv2/quantizer.py +++ b/auto_round/algorithms/quantization/sign_roundv2/quantizer.py @@ -43,7 +43,7 @@ search_optimized_init_scale, ) from auto_round.logger import logger -from auto_round.utils import check_to_quantized, compile_func, get_reciprocal +from auto_round.utils import SUPPORTED_LAYER_TYPES, check_to_quantized, compile_func, get_reciprocal from auto_round.wrapper import WrapperLinear, wrapper_block @@ -416,6 +416,6 @@ def collect_imatrix(module, input, output): handles = [] for _, module in model.named_modules(): - if check_to_quantized(module): + if isinstance(module, SUPPORTED_LAYER_TYPES) and check_to_quantized(module): handles.append(module.register_forward_hook(collect_imatrix)) return handles