Remove preprocess_name parameter from match_named_modules

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

Author: fynnsu

src/compressed_tensors/quantization/lifecycle/apply.py

@@ -38,8 +38,8 @@
     infer_quantization_status,
     is_kv_cache_quant_scheme,
 )
-from compressed_tensors.utils.helpers import fix_fsdp_module_name, replace_module
-from compressed_tensors.utils.match import is_match, match_named_modules, match_targets
+from compressed_tensors.utils.helpers import replace_module
+from compressed_tensors.utils.match import 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
@@ -146,11 +146,7 @@ def apply_quantization_config(
 
     # mark appropriate layers for quantization by setting their quantization schemes
     for name, submodule in match_named_modules(
-        model,
-        target_to_scheme,
-        config.ignore or [],
-        warn_on_fail=True,
-        preprocess_name=fix_fsdp_module_name,
+        model, target_to_scheme, config.ignore or [], warn_on_fail=True
     ):
         # mark modules to be quantized by adding
         # quant scheme to the matching layers

src/compressed_tensors/utils/match.py

@@ -15,7 +15,7 @@
 import logging
 import re
 from collections.abc import Generator
-from typing import Callable, Iterable, List, Tuple
+from typing import Iterable, List, Tuple
 
 import torch
 from compressed_tensors.utils.internal import InternalModule
@@ -37,8 +37,7 @@ def match_named_modules(
     targets: Iterable[str] | None,
     ignore: Iterable[str] | None = None,
     warn_on_fail: bool = False,
-    preprocess_name: Callable[[str], str] = lambda x: x,
-) -> Generator[Tuple[str, torch.nn.Module] | Tuple[str, torch.nn.Module, List[str]]]:
+) -> Generator[Tuple[str, torch.nn.Module]]:
     """
     Yields names and modules which match `targets` but do not match `ignore`.
     Values are returned in order of `model.named_modules()`
@@ -59,9 +58,6 @@ def match_named_modules(
         if isinstance(module, InternalModule):
             continue
 
-        # preprocess the module name and module
-        name = preprocess_name(name)
-
         if any(is_match(name, module, ign) for ign in ignore):
             continue