Skip to content

Commit b1fa4df

Browse files
committed
Remove preprocess_name parameter from match_named_modules
Signed-off-by: Fynn Schmitt-Ulms <[email protected]>
1 parent d94f655 commit b1fa4df

File tree

2 files changed

+5
-13
lines changed

2 files changed

+5
-13
lines changed

src/compressed_tensors/quantization/lifecycle/apply.py

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -38,8 +38,8 @@
3838
infer_quantization_status,
3939
is_kv_cache_quant_scheme,
4040
)
41-
from compressed_tensors.utils.helpers import fix_fsdp_module_name, replace_module
42-
from compressed_tensors.utils.match import is_match, match_named_modules, match_targets
41+
from compressed_tensors.utils.helpers import replace_module
42+
from compressed_tensors.utils.match import match_named_modules, match_targets
4343
from compressed_tensors.utils.offload import update_parameter_data
4444
from compressed_tensors.utils.safetensors_load import get_safetensors_folder
4545
from safetensors import safe_open
@@ -146,11 +146,7 @@ def apply_quantization_config(
146146

147147
# mark appropriate layers for quantization by setting their quantization schemes
148148
for name, submodule in match_named_modules(
149-
model,
150-
target_to_scheme,
151-
config.ignore or [],
152-
warn_on_fail=True,
153-
preprocess_name=fix_fsdp_module_name,
149+
model, target_to_scheme, config.ignore or [], warn_on_fail=True
154150
):
155151
# mark modules to be quantized by adding
156152
# quant scheme to the matching layers

src/compressed_tensors/utils/match.py

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
import logging
1616
import re
1717
from collections.abc import Generator
18-
from typing import Callable, Iterable, List, Tuple
18+
from typing import Iterable, List, Tuple
1919

2020
import torch
2121
from compressed_tensors.utils.internal import InternalModule
@@ -37,8 +37,7 @@ def match_named_modules(
3737
targets: Iterable[str] | None,
3838
ignore: Iterable[str] | None = None,
3939
warn_on_fail: bool = False,
40-
preprocess_name: Callable[[str], str] = lambda x: x,
41-
) -> Generator[Tuple[str, torch.nn.Module] | Tuple[str, torch.nn.Module, List[str]]]:
40+
) -> Generator[Tuple[str, torch.nn.Module]]:
4241
"""
4342
Yields names and modules which match `targets` but do not match `ignore`.
4443
Values are returned in order of `model.named_modules()`
@@ -59,9 +58,6 @@ def match_named_modules(
5958
if isinstance(module, InternalModule):
6059
continue
6160

62-
# preprocess the module name and module
63-
name = preprocess_name(name)
64-
6561
if any(is_match(name, module, ign) for ign in ignore):
6662
continue
6763

0 commit comments

Comments
 (0)