diff --git a/src/compressed_tensors/utils/match.py b/src/compressed_tensors/utils/match.py index abe14f988..4ac9f1ebd 100644 --- a/src/compressed_tensors/utils/match.py +++ b/src/compressed_tensors/utils/match.py @@ -31,6 +31,14 @@ FusedMappping = Mapping[str, Iterable[str]] +def _target_priority_key(target: str) -> tuple: + """Sort key for match_targets: exact names, then longer regexes.""" + if target.startswith("re:"): + pattern = target.removeprefix("re:") + return (1, -len(pattern), pattern) + return (0, 0, target) + + def match_named_modules( model: torch.nn.Module, targets: Iterable[str] | None, @@ -135,10 +143,10 @@ def match_targets( # specific to least specific, and this order will be used when merging configs. # The entries are sorted in the following order: # 1. matches on exact strings - # 2. matches on regex patterns + # 2. matches on regex patterns (longer patterns before shorter ones) # 3. matches on module names (e.g. "Linear") - targets = sorted(targets, key=lambda x: ("re:" in x, x)) + targets = sorted(targets, key=_target_priority_key) matched_targets = [] for target in targets: if match_name(name, target): diff --git a/tests/test_utils/test_match.py b/tests/test_utils/test_match.py index faec318db..2bd2c1564 100644 --- a/tests/test_utils/test_match.py +++ b/tests/test_utils/test_match.py @@ -18,6 +18,7 @@ match_named_modules, match_named_parameters, match_quantizable_tensors, + match_targets, ) from compressed_tensors.utils.match import _match_class from transformers import AutoModelForCausalLM @@ -233,6 +234,35 @@ def test_fused_mapping(self): assert is_match("dummy.gate_up_proj", linear, "Linear", fused=mapping) +class TestMatchTargets: + """Test cases for match_targets priority ordering""" + + def test_regex_specificity_beats_lexicographic_order(self): + linear = nn.Linear(10, 20) + name = "transformer.layers.0.self_attn.q_proj" + matches = match_targets(name, linear, ["re:.*proj", "re:.*q_proj"]) + + assert matches[0] == "re:.*q_proj" + assert "re:.*proj" in matches + + def test_exact_name_beats_regex(self): + linear = nn.Linear(10, 20) + name = "model.layers.0.mlp.down_proj" + matches = match_targets( + name, linear, ["re:.*down_proj", "model.layers.0.mlp.down_proj"] + ) + + assert matches[0] == "model.layers.0.mlp.down_proj" + assert "re:.*down_proj" in matches + + def test_class_match_comes_after_name_matches(self): + linear = nn.Linear(10, 20) + matches = match_targets("layer1", linear, ["Linear", "layer1"]) + + assert matches[0] == "layer1" + assert matches[-1] == "Linear" + + class TestMatchNamedModules: """Test cases for match_named_modules function"""