get_matching_layer -> match_modules_set, match_modules, get_linear_layers ->match_named_modules

Author: ved1beta

src/llmcompressor/modifiers/smoothquant/base.py

@@ -14,9 +14,7 @@
     handle_mapping_resolution_errors,
 )
 from llmcompressor.utils.fsdp.helpers import get_fsdp_parent
-from llmcompressor.utils.pytorch.module import (
-    get_matching_layer,
-)
+from compressed_tensors import match_modules_set
 from compressed_tensors import match_named_modules
 MINIMUM_SMOOTHING_SCALE = 1e-5
 
@@ -208,7 +206,7 @@ def _resolve_mappings(self, model: Module) -> List[SmoothQuantMapping]:
                     balance_layers = []
                     for balance_suffix in to_balance:
                         # find the submodule that matches the activation layer
-                        _, balance_layer = get_matching_layer(
+                        _, balance_layer =match_modules_set(
                             balance_suffix, layer_name, model
                         )
                         if balance_layer:

src/llmcompressor/pipelines/layer_sequential/pipeline.py

@@ -12,10 +12,10 @@
 from llmcompressor.pipelines.cache import IntermediatesCache
 from llmcompressor.pipelines.layer_sequential.helpers import (
     capture_first_layer_intermediates,
-    match_modules,
     maybe_inject_pos_embeddings,
     to_next_layer_kwargs,
 )
+from compressed_tensors import match_named_modules
 from llmcompressor.pipelines.registry import CalibrationPipeline
 from llmcompressor.pipelines.sequential.helpers import (
     dispatch_for_sequential,
@@ -67,7 +67,7 @@ def __call__(
         # find layers
         modifiers = session.lifecycle.recipe.modifiers
         sequential_targets = get_sequential_targets(modifiers, model, dataset_args)
-        layers = match_modules(model, sequential_targets)
+        layers = match_named_modules(model, sequential_targets)
 
         LifecycleCallbacks.calibration_epoch_start()
 

src/llmcompressor/pipelines/sequential/helpers.py

@@ -24,6 +24,7 @@
 from llmcompressor.modifiers.utils.hooks import HooksMixin
 from llmcompressor.utils.helpers import calibration_forward_context, patch_attr
 from llmcompressor.utils.pytorch.module import get_no_split_params
+from compressed_tensors import match_named_modules
 
 from .ast_helpers import autowrap_forwards
 
@@ -100,7 +101,7 @@ def trace_subgraphs(
     :return: a list of Subgraphs in order of execution
     """
     # find modules
-    targets = match_modules(model, sequential_targets)
+    targets = match_named_modules(model, sequential_targets)
     ancestors = get_sequential_ancestors(model, targets)
     offloaded = set(m for m in model.modules() if has_offloaded_params(m))
 

src/llmcompressor/transformers/compression/helpers.py

@@ -8,8 +8,8 @@
 from tqdm import tqdm
 
 from llmcompressor.modifiers import Modifier
-from llmcompressor.pytorch.utils import get_linear_layers
 from llmcompressor.pytorch.utils.helpers import tensor_sparsity
+from compressed_tensors import match_named_modules
 
 __ALL__ = [
     "tensor_follows_mask_structure",
@@ -76,7 +76,7 @@ def infer_sparsity_structure_from_model(model: torch.nn.Module) -> Optional[str]
     # check for the common sparsity structures
     structures = {"2:4"}
     for sparsity_structure in structures:
-        linear_modules = get_linear_layers(model)
+        linear_modules = match_named_modules(model, linear=True)
         offloaded_params = get_state_dict_offloaded_model(model)
 
         linear_modules_with_sparsity_structure = [

tests/llmcompressor/transformers/tracing/test_models.py

@@ -14,9 +14,9 @@
     WhisperForConditionalGeneration,
 )
 
-from llmcompressor.pipelines.sequential.helpers import match_modules
 from llmcompressor.transformers.tracing.debug import trace
 from llmcompressor.utils.pytorch.module import get_no_split_params
+from compressed_tensors import match_named_modules
 
 
 @pytest.mark.skipif(
@@ -148,7 +148,7 @@ def get_target_modules(model, sequential_targets):
     if isinstance(sequential_targets, str):
         sequential_targets = [sequential_targets]
 
-    return match_modules(model, sequential_targets)
+    return match_named_modules(model, sequential_targets)
 
 
 def run_subgraphs(model, subgraphs, inputs):