Skip to content

[Utils] Expand is_match #416

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[Utils] Expand is_match #416

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion src/compressed_tensors/transform/factory/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@

from abc import ABC, abstractmethod
from collections import defaultdict
from typing import List, Optional, Tuple, Set
from typing import List, Optional, Set, Tuple

import torch
import torch.nn.utils.parametrize as P
Expand Down
24 changes: 16 additions & 8 deletions src/compressed_tensors/utils/match.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
import logging
import re
from collections.abc import Generator
from typing import Iterable, Mapping, Optional, Tuple
from typing import Iterable, List, Mapping, Optional, Tuple

import torch
from compressed_tensors.utils.internal import InternalModule
Expand Down Expand Up @@ -57,10 +57,10 @@ def match_named_modules(
unmatched_targets = set(targets)
for name, module in model.named_modules():
for target in targets:
if is_match(name, module, target, fused):
if is_match(name, module, target, fused=fused):
unmatched_targets -= {target}

if not any(is_match(name, module, ign, fused) for ign in ignore):
if not is_match(name, module, ignore, fused=fused):
yield name, module

if warn_on_fail:
Expand Down Expand Up @@ -155,9 +155,7 @@ def match_modules_set(
for name, module in model.named_modules():
# match until we get a full set
for target in targets:
if is_match(name, module, target) and not any(
is_match(name, module, ign) for ign in ignore
):
if is_match(name, module, target, ignore):
if matches[target] is not None:
raise ValueError(f"Matched a {target} twice before completing set")
matches[target] = module
Expand All @@ -176,7 +174,8 @@ def match_modules_set(
def is_match(
name: str,
module: torch.nn.Module,
target: str,
targets: str | Iterable[str],
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is this rename a potential breaking change in vllm?

ignore: str | Iterable[str] = tuple(),
fused: Optional[FusedMappping] = None,
) -> bool:
"""
Expand All @@ -198,8 +197,17 @@ def is_match(
:fused: optional mapping from suffixes of fused modules to the suffixes of their
corresponding shards
"""
targets = [targets] if isinstance(targets, str) else targets
ignore = [ignore] if isinstance(ignore, str) else ignore

return not isinstance(module, InternalModule) and (
_match_name(name, target, fused) or _match_class(module, target)
any(
_match_name(name, target, fused) or _match_class(module, target)
for target in targets
)
and not any(
_match_name(name, ign, fused) or _match_class(module, ign) for ign in ignore
)
)


Expand Down
22 changes: 14 additions & 8 deletions tests/test_utils/test_match.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -201,14 +201,20 @@ def test_fused_mapping(self):
"gate_up_proj": ["gate_proj", "up_proj"],
}

assert is_match("dummy.qkv_proj", linear, "re:.*q_proj", mapping) == True
assert is_match("dummy.qkv_proj", linear, "re:.*k_proj", mapping) == True
assert is_match("dummy.qkv_proj", linear, "re:.*v_proj", mapping) == True
assert is_match("dummy.qkv_proj", linear, "Linear", mapping) == True

assert is_match("dummy.gate_up_proj", linear, "re:.*gate_proj", mapping) == True
assert is_match("dummy.gate_up_proj", linear, "re:.*up_proj", mapping) == True
assert is_match("dummy.gate_up_proj", linear, "Linear", mapping) == True
assert is_match("dummy.qkv_proj", linear, "re:.*q_proj", fused=mapping) == True
assert is_match("dummy.qkv_proj", linear, "re:.*k_proj", fused=mapping) == True
assert is_match("dummy.qkv_proj", linear, "re:.*v_proj", fused=mapping) == True
assert is_match("dummy.qkv_proj", linear, "Linear", fused=mapping) == True

assert (
is_match("dummy.gate_up_proj", linear, "re:.*gate_proj", fused=mapping)
== True
)
assert (
is_match("dummy.gate_up_proj", linear, "re:.*up_proj", fused=mapping)
== True
)
assert is_match("dummy.gate_up_proj", linear, "Linear", fused=mapping) == True


class TestMatchNamedModules:
Expand Down