diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 448f085adb7..da2a422011e 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -27,6 +27,7 @@ Changelog - Fix ONNX preprocessing with custom TensorRT plugin libraries on TensorRT 11. - Fix ``INT8_DEFAULT_CFG``, ``FP8_DEFAULT_CFG``, ``NVFP4_DEFAULT_CFG`` and the other presets and recipes that import ``configs/ptq/units/default_disabled_quantizers`` and enable ``*input_quantizer``: since 0.44 they also quantized ``nn.LayerNorm`` inputs, which adds quantization error to the residual stream of models built on ``nn.LayerNorm`` such as BERT-style encoders. Re-quantize affected models; a recipe that wants ``nn.LayerNorm`` quantizers must enable them after importing that unit. - Fix ONNX quantization with the default ``concat_elimination`` pass crashing in calibration with ``KeyError: '_ReduceMin'`` when a ``Concat`` input is a constant, such as the class token of a ViT. +- Fix ONNX AWQ-Lite's ``get_weight_scale`` computing per-block weight scales from the wrong rows: reshaping the weight directly into ``(block_size, -1)`` interleaves rows from different blocks instead of grouping each block's ``block_size`` consecutive rows together, silently skewing the search scale AWQ-Lite picks for INT4 PTQ. 0.48.0 (2026-10-xx) ^^^^^^^^^^^^^^^^^^^ diff --git a/modelopt/onnx/quantization/int4.py b/modelopt/onnx/quantization/int4.py index fc273d9f4f0..6578d2f9c60 100644 --- a/modelopt/onnx/quantization/int4.py +++ b/modelopt/onnx/quantization/int4.py @@ -747,10 +747,16 @@ def get_weight_scale(weight, block_size=None): slice_after_padding = slice(org_shape[0]) weight = _pad(weight, block_size) org_shape = weight.shape - weight = weight.reshape(block_size, -1) - weight_abs_amax = weight.__abs__().max(axis=0, keepdims=True) - scale = weight.__abs__() / (weight_abs_amax + np.finfo(weight_abs_amax.dtype).tiny) - scale = scale.reshape(org_shape) + # Blocks are `block_size` consecutive rows; splitting dim 0 into + # (n_blocks, block_size) keeps each block's rows together, unlike + # reshape(block_size, -1), which interleaves rows from different blocks. + weight = weight.reshape(org_shape[0] // block_size, block_size, org_shape[1]) + weight_abs_amax = weight.__abs__().max(axis=1, keepdims=True) + scale = weight.__abs__() / (weight_abs_amax + np.finfo(weight_abs_amax.dtype).tiny) + scale = scale.reshape(org_shape) + else: + weight_abs_amax = weight.__abs__().max(axis=0, keepdims=True) + scale = weight.__abs__() / (weight_abs_amax + np.finfo(weight_abs_amax.dtype).tiny) if slice_after_padding is not None: scale = scale[slice_after_padding, ...] scale = scale.mean(1) diff --git a/tests/unit/onnx/quantization/test_int4.py b/tests/unit/onnx/quantization/test_int4.py new file mode 100644 index 00000000000..0403f654aea --- /dev/null +++ b/tests/unit/onnx/quantization/test_int4.py @@ -0,0 +1,84 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np + +from modelopt.onnx.quantization.int4 import get_weight_scale + + +def _reference_weight_scale(weight: np.ndarray, block_size: int) -> np.ndarray: + """Per-block-max scale using consecutive, non-overlapping row blocks.""" + ci, co = weight.shape + n_blocks = -(-ci // block_size) # ceil division, matches the padding in get_weight_scale + padded = np.pad(weight, ((0, n_blocks * block_size - ci), (0, 0))) + scale = np.zeros_like(padded) + for b in range(n_blocks): + block = padded[b * block_size : (b + 1) * block_size] + amax = np.abs(block).max(axis=0, keepdims=True) + scale[b * block_size : (b + 1) * block_size] = np.abs(block) / ( + amax + np.finfo(amax.dtype).tiny + ) + return scale[:ci].mean(1) + + +def test_get_weight_scale_groups_consecutive_rows_into_blocks(): + """A block's max must come from its own rows, not rows from a different block. + + Reshaping (ci, co) directly into (block_size, -1) does not slice `ci` into + contiguous blocks of `block_size` rows -- it interleaves rows from every block, + so a row's scale ends up computed against the wrong max. + """ + weight = np.zeros((8, 1)) + weight[0, 0] = 10.0 + weight[1, 0] = 5.0 + + scale = get_weight_scale(weight, block_size=4) + + # Rows 0 and 1 are both in the first 4-row block, whose max is 10. + assert scale[0] == 1.0 + assert scale[1] == 0.5 + np.testing.assert_array_equal(scale[2:], 0.0) + + +def test_get_weight_scale_matches_reference_when_divisible(): + rng = np.random.RandomState(0) + weight = rng.randn(16, 3) + block_size = 4 + + scale = get_weight_scale(weight.copy(), block_size) + + np.testing.assert_allclose(scale, _reference_weight_scale(weight, block_size)) + + +def test_get_weight_scale_matches_reference_with_padding(): + rng = np.random.RandomState(1) + weight = rng.randn(10, 3) # not a multiple of block_size + block_size = 4 + + scale = get_weight_scale(weight.copy(), block_size) + + assert scale.shape == (10,) + np.testing.assert_allclose(scale, _reference_weight_scale(weight, block_size)) + + +def test_get_weight_scale_without_block_size_is_unchanged(): + rng = np.random.RandomState(2) + weight = rng.randn(16, 3) + + scale = get_weight_scale(weight.copy(), block_size=None) + + amax = np.abs(weight).max(axis=0, keepdims=True) + expected = (np.abs(weight) / (amax + np.finfo(amax.dtype).tiny)).mean(1) + np.testing.assert_allclose(scale, expected)