From 85f82d9d3ba86d291294e6afd00b17b771e27772 Mon Sep 17 00:00:00 2001 From: arham766 Date: Sun, 5 Jul 2026 13:39:34 -0700 Subject: [PATCH] fix: don't normalize SequentialQuantizer child names in wildcard matching MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The fused-experts name normalization from #1340 collapses weight_quantizer. to weight_quantizer — intended for per-expert nn.ModuleList layouts, but this is exactly the child naming of a SequentialQuantizer stored at attribute weight_quantizer. Its children are TensorQuantizers, so they matched singular wildcards like *weight_quantizer on their own, with two user-visible consequences: - applying the same list config twice via set_quantizer_attributes_full nested a SequentialQuantizer inside each sub-slot (non-idempotent, corrupt state no consumer supports); - set_quantizer_attributes_partial with a list config and a wildcard raised ValueError mid-iteration for targets that are already SequentialQuantizers, contradicting its own docstring and leaving the model partially configured. Skip the normalized match when the quantizer's direct parent is a SequentialQuantizer. Raw wildcard matches (including explicit *weight_quantizer.0), callable filters, and the plural weight_quantizers.N fused-experts path are unchanged; the parent lookup only runs when the raw match already failed. Part of the findings in #1902. Signed-off-by: arham766 --- modelopt/torch/quantization/conversion.py | 23 ++++-- .../torch/quantization/test_quantize_cpu.py | 71 ++++++++++++++++++- 2 files changed, 87 insertions(+), 7 deletions(-) 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."""