Simplify signature of match_named_modules

Removed `yield_matched_targets` and `warn_on_unmatched_ignores` and updated rest of code

Signed-off-by: Fynn Schmitt-Ulms <[email protected]>

Author: fynnsu

src/compressed_tensors/quantization/lifecycle/apply.py

@@ -39,7 +39,7 @@
     is_kv_cache_quant_scheme,
 )
 from compressed_tensors.utils.helpers import fix_fsdp_module_name, replace_module
-from compressed_tensors.utils.match import match_named_modules
+from compressed_tensors.utils.match import is_match, match_named_modules, match_targets
 from compressed_tensors.utils.offload import update_parameter_data
 from compressed_tensors.utils.safetensors_load import get_safetensors_folder
 from safetensors import safe_open
@@ -145,17 +145,16 @@ def apply_quantization_config(
         from compressed_tensors.linear.compressed_linear import CompressedLinear
 
     # mark appropriate layers for quantization by setting their quantization schemes
-    for name, submodule, matched_targets in match_named_modules(
+    for name, submodule in match_named_modules(
         model,
         target_to_scheme,
         config.ignore or [],
         warn_on_fail=True,
-        warn_on_unmatched_ignores=True,
-        yield_matched_targets=True,
         preprocess_name=fix_fsdp_module_name,
     ):
         # mark modules to be quantized by adding
         # quant scheme to the matching layers
+        matched_targets = list(match_targets(name, submodule, target_to_scheme))
         scheme = _scheme_from_targets(target_to_scheme, matched_targets, name)
         if run_compressed:
             format = config.format

src/compressed_tensors/utils/match.py

@@ -37,8 +37,6 @@ def match_named_modules(
     targets: Iterable[str] | None,
     ignore: Iterable[str] | None = None,
     warn_on_fail: bool = False,
-    warn_on_unmatched_ignores: bool = False,
-    yield_matched_targets: bool = False,
     preprocess_name: Callable[[str], str] = lambda x: x,
 ) -> Generator[Tuple[str, torch.nn.Module] | Tuple[str, torch.nn.Module, List[str]]]:
     """
@@ -49,70 +47,36 @@ def match_named_modules(
     :param targets: target strings, potentially containing "re:" prefixes
     :param ignore: targets to ignore, potentially containing "re:" prefixes
     :param warn_on_fail: if True, warns if any targets do not match any modules in model
-    :param warn_on_unmatched_ignores: if True, warns if any ignores do not match any modules in model
-    :param yield_matched_targets: if True, yields the matched targets in addition to the module name and module
     :param preprocess_name: a function to preprocess the module name
     :return: generator of module names and modules
     """
     ignore = ignore or []
     targets = targets or []
 
     unmatched_targets = set(targets)
-    unmatched_ignores = set(ignore)
 
-    # Note: when yield_matched_targets is True, the ordering of the targets is important
-    # Order targets by type: exact name match, regex name match, class name match
-    targets = sorted(targets, key=lambda x: ("re:" in x, x))
     for name, module in model.named_modules():
         if isinstance(module, InternalModule):
             continue
 
         # preprocess the module name and module
         name = preprocess_name(name)
 
-        ignore_matched = False
-        for ign in ignore:
-            if is_match(name, module, ign):
-                unmatched_ignores -= {ign}
-                ignore_matched = True
-                break
-        if ignore_matched:
+        if any(is_match(name, module, ign) for ign in ignore):
             continue
 
-        matched_target_on_name = []
-        matched_target_on_class = []
-        # Check for name matches first (exact then regex, enforced by sort above)
         for target in targets:
-            if _match_name(name, target):
+            if is_match(name, module, target):
                 unmatched_targets -= {target}
-                matched_target_on_name.append(target)
-                if not yield_matched_targets:
-                    break
-            elif _match_class(module, target):
-                unmatched_targets -= {target}
-                matched_target_on_class.append(target)
-                if not yield_matched_targets:
-                    break
-
-        matched_targets = matched_target_on_name + matched_target_on_class
-        if matched_targets:
-            if yield_matched_targets:
-                yield name, module, matched_targets
-            else:
                 yield name, module
+                break
 
     if warn_on_fail:
         for target in unmatched_targets:
             _LOGGER.warning(
                 f"Could not match `{target}` in instance of {model.__class__.__name__}"
             )
 
-    if warn_on_unmatched_ignores:
-        for ign in unmatched_ignores:
-            _LOGGER.warning(
-                f"Unmatched ignore targets: {unmatched_ignores}, in instance of {model.__class__.__name__}"
-            )
-
 
 def match_named_parameters(
     model: torch.nn.Module,
@@ -151,6 +115,23 @@ def match_named_parameters(
             )
 
 
+def match_targets(
+    name: str, module: torch.nn.Module, targets: Iterable[str]
+) -> Generator[str]:
+    """
+    Yields the targets that match the given name and module.
+    Outputs are ordered by type: exact name match, regex name match, class name match
+    """
+    targets = sorted(targets, key=lambda x: ("re:" in x, x))
+    for target in targets:
+        if _match_name(name, target):
+            yield target
+
+    for target in targets:
+        if _match_class(module, target):
+            yield target
+
+
 def match_modules_set(
     model: torch.nn.Module,
     targets: Iterable[str],

tests/test_quantization/lifecycle/test_apply.py

@@ -258,15 +258,13 @@ def get_sample_tinyllama_quant_config(status: str = "frozen"):
 
 @requires_accelerate()
 @pytest.mark.parametrize(
-    "ignore,should_raise_warning",
+    "ignore",
     [
-        [("lm_head", "re:.*gate"), False],
-        [("lm_head", "re:.*foobarbaz"), True],
+        ("lm_head", "re:.*gate"),
+        ("lm_head", "re:.*foobarbaz"),
     ],
 )
-def test_apply_quantization_status(caplog, ignore, should_raise_warning):
-    import logging
-
+def test_apply_quantization_status(ignore):
     # load a dense, unquantized tiny llama model
     model = get_tinyllama_model()
     quantization_config_dict = {
@@ -290,10 +288,4 @@ def test_apply_quantization_status(caplog, ignore, should_raise_warning):
     config = QuantizationConfig(**quantization_config_dict)
     config.quantization_status = QuantizationStatus.CALIBRATION
 
-    # mismatch in the ignore key of quantization_config_dict
-    with caplog.at_level(logging.WARNING):
-        apply_quantization_config(model, config)
-        if should_raise_warning:
-            assert len(caplog.text) > 0
-        else:
-            assert len(caplog.text) == 0
+    apply_quantization_config(model, config)