diff --git a/modelopt/torch/quantization/conversion.py b/modelopt/torch/quantization/conversion.py index fa7a8a0a128..8bb76a4985e 100644 --- a/modelopt/torch/quantization/conversion.py +++ b/modelopt/torch/quantization/conversion.py @@ -348,12 +348,23 @@ def _match_quantizer( if not isinstance(module, (TensorQuantizer, SequentialQuantizer)): return False if isinstance(wildcard_or_filter_func, str): - normalized = _normalize_fused_experts_quantizer_name(name) - if not ( - fnmatch.fnmatch(name, wildcard_or_filter_func) - or (normalized != name and fnmatch.fnmatch(normalized, wildcard_or_filter_func)) - ): - return False + if not fnmatch.fnmatch(name, wildcard_or_filter_func): + normalized = _normalize_fused_experts_quantizer_name(name) + if normalized == name or not fnmatch.fnmatch(normalized, wildcard_or_filter_func): + return False + # The index-stripping normalization exists solely for fused-experts layouts + # that hold per-expert quantizers in an ``nn.ModuleList``. A + # ``SequentialQuantizer`` stored at attribute ``weight_quantizer`` exposes + # children with the same dotted shape (``...weight_quantizer.0``), so + # normalizing them would make each sub-quantizer match singular wildcards + # like ``*weight_quantizer`` on its own. That corrupts re-application of + # list configs (nesting SequentialQuantizers inside sub-slots) and breaks + # list-based partial updates on existing SequentialQuantizers. Sub-quantizers + # must only be addressed through their container, so skip the normalized + # match when the quantizer's direct parent is a SequentialQuantizer. + parent_name = name.rpartition(".")[0] + if isinstance(full_model.get_submodule(parent_name), SequentialQuantizer): + return False elif callable(wildcard_or_filter_func): if not wildcard_or_filter_func(name): return False diff --git a/tests/unit/torch/quantization/test_quantize_cpu.py b/tests/unit/torch/quantization/test_quantize_cpu.py index 3e4925e7b63..ae3899f6938 100644 --- a/tests/unit/torch/quantization/test_quantize_cpu.py +++ b/tests/unit/torch/quantization/test_quantize_cpu.py @@ -33,7 +33,10 @@ import modelopt.torch.quantization as mtq from modelopt.torch.quantization.calib import MaxCalibrator from modelopt.torch.quantization.config import QuantizerAttributeConfig -from modelopt.torch.quantization.conversion import set_quantizer_attributes_full +from modelopt.torch.quantization.conversion import ( + set_quantizer_attributes_full, + set_quantizer_attributes_partial, +) from modelopt.torch.quantization.nn.modules.tensor_quantizer import ( SequentialQuantizer, TensorQuantizer, @@ -401,6 +404,72 @@ def test_list_attributes_creates_sequential_quantizer(self): assert isinstance(module, SequentialQuantizer) assert len(module) == 2 + def test_list_attributes_double_application_is_idempotent(self): + """Applying the same list config twice must not nest SequentialQuantizers. + + Regression test: the fused-experts name normalization used to collapse a + SequentialQuantizer child name like ``...weight_quantizer.0`` down to + ``...weight_quantizer``, so on a second pass each sub-quantizer matched the + singular wildcard on its own and was replaced with a nested + SequentialQuantizer. + """ + model = self._quantize(SimpleLinear()) + attrs = [ + QuantizerAttributeConfig(num_bits=4, block_sizes={-1: 128}), + QuantizerAttributeConfig(num_bits=8, axis=0), + ] + set_quantizer_attributes_full(model, "*weight_quantizer", attrs) + set_quantizer_attributes_full(model, "*weight_quantizer", attrs) + seen = 0 + for name, module in model.named_modules(): + if name.endswith("weight_quantizer"): + seen += 1 + assert isinstance(module, SequentialQuantizer) + assert len(module) == 2 + for sub in module: + assert isinstance(sub, TensorQuantizer) + assert not isinstance(sub, SequentialQuantizer), ( + f"{name} has a nested SequentialQuantizer sub-quantizer" + ) + assert module[0].num_bits == 4 + assert module[1].num_bits == 8 + assert seen > 0 + + +def test_partial_list_on_existing_sequential_quantizer(): + """A list partial config applies per-position to an existing SequentialQuantizer. + + This is the documented contract of set_quantizer_attributes_partial. Regression + test: the fused-experts name normalization made the SequentialQuantizer's child + TensorQuantizers (``...weight_quantizer.0``) match ``*weight_quantizer`` too, and + the list-on-TensorQuantizer check then raised ValueError mid-iteration. + """ + model = mtq.quantize(SimpleLinear(), mtq.INT8_DEFAULT_CFG, lambda m: m(m.get_input())) + set_quantizer_attributes_full( + model, + "*weight_quantizer", + [ + QuantizerAttributeConfig(num_bits=4, block_sizes={-1: 128}), + QuantizerAttributeConfig(num_bits=8, axis=0), + ], + ) + # Must not raise: only the SequentialQuantizer containers should match, never + # their child TensorQuantizers. + set_quantizer_attributes_partial( + model, "*weight_quantizer", [{"enable": False}, {"num_bits": 4}] + ) + seen = 0 + for name, module in model.named_modules(): + if name.endswith("weight_quantizer"): + seen += 1 + assert isinstance(module, SequentialQuantizer) + assert not module[0].is_enabled, f"{name}[0] should be disabled" + assert module[0].num_bits == 4, "unspecified attribute should be preserved" + assert module[1].is_enabled, f"{name}[1] should stay enabled" + assert module[1].num_bits == 4, "per-position update should apply to slot 1" + assert module[1].axis == 0, "unspecified attribute should be preserved" + assert seen > 0 + def test_ordering_later_entry_overrides_earlier(): """Later entries in quant_cfg override earlier ones for the same quantizer."""