removed get_layer, get_layers -> match_named_modules

Author: ved1beta

src/llmcompressor/modifiers/awq/base.py

@@ -30,7 +30,7 @@
 from llmcompressor.pipelines.cache import IntermediatesCache
 from llmcompressor.utils.fsdp.helpers import get_fsdp_parent
 from llmcompressor.utils.helpers import calibration_forward_context
-from llmcompressor.utils.pytorch.module import get_layers
+from compressed_tensors import match_named_modules
 
 __all__ = ["AWQModifier"]
 
@@ -305,7 +305,7 @@ def _set_resolved_mappings(self, model: Module) -> None:
         """
         resolved_mappings: list[ResolvedMapping] = []
         for mapping_idx, mapping in enumerate(self.mappings):
-            smooth_layers = get_layers(
+            smooth_layers = match_named_modules(
                 mapping.smooth_layer, model, exclude_internal_modules=True
             )
             smooth_names = [
@@ -329,7 +329,7 @@ def _set_resolved_mappings(self, model: Module) -> None:
                 balance_layers, balance_names = [], []
                 for balance_regex in mapping.balance_layers:
                     # find the submodules that match the activation layer
-                    for balance_suffix, balance_layer in get_layers(
+                    for balance_suffix, balance_layer in match_named_modules(
                         balance_regex,
                         smooth_parent,
                         exclude_internal_modules=True,

src/llmcompressor/modifiers/distillation/output/base.py

@@ -11,7 +11,7 @@
 )
 from llmcompressor.utils.fsdp.context import summon_full_params_context
 from llmcompressor.utils.fsdp.helpers import maybe_get_wrapped, set_wrapped_model
-from llmcompressor.utils.pytorch.module import get_layers
+from compressed_tensors import match_named_modules 
 
 __all__ = ["OutputDistillationModifier"]
 
@@ -61,8 +61,8 @@ def on_initialize(self, state: State, **kwargs) -> bool:
             else:
                 model_target, teacher_target = target, target
 
-            model_layers = get_layers(model_target, state.model)
-            teacher_layers = get_layers(teacher_target, state.teacher_model)
+            model_layers = match_named_modules(model_target, state.model)
+            teacher_layers = match_named_modules(teacher_target, state.teacher_model)
 
             if len(model_layers) < 1:
                 raise ValueError(f"no model layers found for target {target}")

src/llmcompressor/modifiers/obcq/sgpt_base.py

@@ -13,11 +13,11 @@
 from llmcompressor.modifiers.modifier import Modifier
 from llmcompressor.modifiers.utils.hooks import HooksMixin
 from llmcompressor.utils.pytorch.module import (
-    get_layers,
     get_no_split_params,
     get_prunable_layers,
     match_targets,
 )
+from compressed_tensors import match_named_modules
 
 
 class SparsityModifierBase(Modifier):
@@ -114,8 +114,8 @@ def on_initialize(self, state: "State", **kwargs) -> bool:
 
         # infer module and sequential targets
         self.sequential_targets = self._infer_sequential_targets(model)
-        layers = get_layers(self.sequential_targets, model)
-        self._target_layers = get_layers(
+        layers = match_named_modules(self.sequential_targets, model)
+        self._target_layers = match_named_modules(
             self.targets, model
         )  # layers containing targets
 

src/llmcompressor/modifiers/smoothquant/base.py

@@ -15,11 +15,10 @@
 )
 from llmcompressor.utils.fsdp.helpers import get_fsdp_parent
 from llmcompressor.utils.pytorch.module import (
-    get_layers,
     get_matching_layer,
     match_targets,
 )
-
+from compressed_tensors import match_named_modules
 MINIMUM_SMOOTHING_SCALE = 1e-5
 
 
@@ -204,7 +203,7 @@ def _resolve_mappings(self, model: Module) -> List[SmoothQuantMapping]:
         """
         resolved_mappings = []
         for to_balance, to_smooth in self.mappings:
-            to_smooth_layers = get_layers(to_smooth, model)
+            to_smooth_layers = match_named_modules(to_smooth, model)
             for layer_name, smooth_layer in to_smooth_layers.items():
                 if not match_targets(layer_name, self.ignore)[0]:
                     balance_layers = []

src/llmcompressor/utils/pytorch/module.py

@@ -19,6 +19,7 @@
     fix_fsdp_module_name,
     summon_full_params_context,
 )
+from compressed_tensors import match_named_modules
 
 try:
     quant_err = None
@@ -49,8 +50,6 @@
     "match_targets",
     "get_default_params",
     "match_layers_params",
-    "get_layers",
-    "get_layer",
     "get_terminal_layers",
     "get_prunable_layers",
     "get_quantizable_layers",
@@ -155,45 +154,6 @@ def match_layers_params(
     return resolved
 
 
-def get_layers(
-    targets: Union[str, List[str]],
-    module: Module,
-    exclude_internal_modules: bool = False,
-) -> Dict[str, Module]:
-    """
-    Get layers (also known as submodules) of module based on targets
-
-    :param targets: names or regexes to search for
-        Can be regex, e.g. "re:.*input_layernorm$" to find all layers
-        in module whose names end in string "input_layernorm"
-    :param module: Parent module in which to search for targets
-    :param exclude_internal_modules: If True, don't include internal
-        modules added by llm-compressor, e.g. Observers and Transforms.
-        Defaults to False to maintain backward compatibility
-
-    :return: dict of {layer name -> module} of all layers in module
-        that match targets
-    """
-    layer_dict = match_layers_params(targets, module)
-    if exclude_internal_modules:
-        layer_dict = {
-            name: layer
-            for name, layer in layer_dict.items()
-            if not isinstance(layer, InternalModule)
-        }
-
-    return layer_dict
-
-
-def get_layer(target: str, module: Module) -> Tuple[str, Module]:
-    layers = get_layers(target, module)
-    if len(layers) != 1:
-        raise ValueError(f"Expected 1 layer for target {target}, found {len(layers)}")
-    name, layer = next(iter(layers.items()))
-
-    return name, layer
-
-
 def get_terminal_layers(module: Module) -> Dict[str, Module]:
     terminal = {}
 
@@ -271,7 +231,7 @@ def get_matching_layer(
     :return: Tuple containing the layer name and module that fits the target regex and
     best matches name_to_match, or None if no match can be found
     """
-    potential_matches = get_layers(target, module)
+    potential_matches = match_named_modules(target, module)
     largest_substring = 0
     match = None
     for name, module in potential_matches.items():

tests/llmcompressor/transformers/obcq/test_obcq_owl.py

@@ -6,8 +6,7 @@
 from llmcompressor.core.session_functions import create_session
 from llmcompressor.datasets import format_calibration_data
 from llmcompressor.modifiers.obcq import SparseGPTModifier
-from llmcompressor.utils.pytorch.module import get_layers
-
+from compressed_tensors import  match_named_modules
 
 @pytest.mark.integration
 def test_infer_owl_layer_sparsity():
@@ -29,7 +28,7 @@ def test_infer_owl_layer_sparsity():
         dataloader = format_calibration_data(dataset)
 
         sequential_targets = modifier._infer_sequential_targets(model)
-        layers = get_layers(sequential_targets, model)
+        layers = match_named_modules(sequential_targets, model)
         sparsities = modifier._infer_owl_layer_sparsity(model, layers, dataloader)
         assert sparsities.keys() == layers.keys()