removed get_quantizable,prunable,terminal_layers

Author: ved1beta

src/llmcompressor/modifiers/obcq/sgpt_base.py

@@ -14,8 +14,6 @@
 from llmcompressor.modifiers.utils.hooks import HooksMixin
 from llmcompressor.utils.pytorch.module import (
     get_no_split_params,
-    get_prunable_layers,
-    match_targets,
 )
 from compressed_tensors import match_named_modules
 
@@ -149,11 +147,11 @@ def on_start(self, state: State, event: Event, **kwargs):
                 layer_sparsity = self.sparsity[index]
             else:
                 layer_sparsity = self.sparsity
-
-            for name, module in get_prunable_layers(layer).items():
+            prunable_targets = ["Linear", "Conv1d", "Conv2d", "Conv3d", "QATLinear", "QATConv2d", "QATConv3d", "Conv1D"]
+            for name, module in match_named_modules(layer, prunable_targets).items():
                 name = f"{layer_name}.{name}"
 
-                if match_targets(name, self.ignore)[0]:
+                if match_named_modules(name, self.ignore)[0]:
                     continue
 
                 # HACK: previously, embeddings were not quantized because they were not
@@ -210,7 +208,8 @@ def _infer_owl_layer_sparsity(
 
         groups = {}
         for name, layer in layers.items():
-            prunable_layers = get_prunable_layers(layer)
+            prunable_targets = ["Linear", "Conv1d", "Conv2d", "Conv3d", "QATLinear", "QATConv2d", "QATConv3d", "Conv1D"]
+            prunable_layers = match_named_modules(layer, prunable_targets)
             z = [
                 m.weight.abs() * activations[f"{name}.{n}"].unsqueeze(0)
                 for n, m in prunable_layers.items()

src/llmcompressor/modifiers/smoothquant/base.py

@@ -16,7 +16,6 @@
 from llmcompressor.utils.fsdp.helpers import get_fsdp_parent
 from llmcompressor.utils.pytorch.module import (
     get_matching_layer,
-    match_targets,
 )
 from compressed_tensors import match_named_modules
 MINIMUM_SMOOTHING_SCALE = 1e-5
@@ -205,7 +204,7 @@ def _resolve_mappings(self, model: Module) -> List[SmoothQuantMapping]:
         for to_balance, to_smooth in self.mappings:
             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]:
+                if not match_named_modules(layer_name, self.ignore)[0]:
                     balance_layers = []
                     for balance_suffix in to_balance:
                         # find the submodule that matches the activation layer

src/llmcompressor/utils/pytorch/module.py

@@ -47,10 +47,6 @@
 
 
 __all__ = [
-    "match_targets",
-    "get_terminal_layers",
-    "get_prunable_layers",
-    "get_quantizable_layers",
     "qat_active",
     "get_matching_layer",
     "get_no_split_params",
@@ -61,21 +57,6 @@
 ALL_QUANTIZABLE_TARGET = "__ALL_QUANTIZABLE__"
 
 
-def match_targets(name: str, targets: Union[str, List[str]]) -> Tuple[bool, int]:
-    if isinstance(targets, str):
-        targets = [targets]
-
-    for index, target in enumerate(targets):
-        if target[:3] == "re:":
-            pattern = target[3:]
-            if re.match(pattern, name):
-                return True, index
-        elif name == target:
-            return True, index
-
-    return False, -1
-
-
 def match_class(layer: Module, targets: Union[str, List[str]]) -> Tuple[bool, int]:
     if isinstance(targets, str):
         targets = [targets]
@@ -87,51 +68,6 @@ def match_class(layer: Module, targets: Union[str, List[str]]) -> Tuple[bool, in
     return False, -1
 
 
-def get_terminal_layers(module: Module) -> Dict[str, Module]:
-    terminal = {}
-
-    for name, layer in module.named_modules():
-        if len(list(layer.named_modules())) > 1:
-            continue
-
-        terminal[name] = layer
-
-    return terminal
-
-
-def get_prunable_layers(module: Module) -> Dict[str, Module]:
-    prunable = {}
-
-    for name, layer in module.named_modules():
-        if (
-            isinstance(layer, Linear)
-            or isinstance(layer, _ConvNd)
-            or (QATLinear and isinstance(layer, QATLinear))
-            or (QATConv2d and isinstance(layer, QATConv2d))
-            or (QATConv3d and isinstance(layer, QATConv3d))
-            or (TransformerConv1D and isinstance(layer, TransformerConv1D))
-        ):
-            prunable[name] = layer
-
-    return prunable
-
-
-def get_quantizable_layers(module: Module) -> Dict[str, Module]:
-    if QATLinear is None:
-        raise ImportError(
-            "PyTorch version is not setup for Quantization. "
-            "Please install a QAT compatible version of PyTorch"
-        )
-
-    quantizable = {}
-
-    for name, layer in module.named_modules():
-        if isinstance(layer, Linear) or isinstance(layer, _ConvNd):
-            quantizable[name] = layer
-
-    return quantizable
-
-
 def qat_active(module: Module) -> bool:
     """
     Determines if any layers in the model have quantization enabled by checking for