required changes

Author: ved1beta

src/llmcompressor/modifiers/awq/base.py

@@ -324,7 +324,7 @@ def _set_resolved_mappings(self, model: Module) -> None:
                 smooth_layer = smooth_layers[smooth_name]
 
                 smooth_parent_name = ".".join(smooth_name.split(".")[:-1])
-                smooth_parent = attrgetter(smooth_parent_name)(model)
+                smooth_parent = attrgetter(smooth_parent_name)(model) if smooth_parent_name else model
 
                 balance_layers, balance_names = [], []
                 for balance_regex in mapping.balance_layers:

src/llmcompressor/modifiers/obcq/sgpt_base.py

@@ -17,6 +17,9 @@
 )
 from compressed_tensors import match_named_modules
 
+def get_prunable_targets():
+    """Return the list of prunable layer types."""
+    return ["Linear", "Conv1d", "Conv2d", "Conv3d", "QATLinear", "QATConv2d", "QATConv3d", "Conv1D"]
 
 class SparsityModifierBase(Modifier):
     """
@@ -147,7 +150,7 @@ def on_start(self, state: State, event: Event, **kwargs):
                 layer_sparsity = self.sparsity[index]
             else:
                 layer_sparsity = self.sparsity
-            prunable_targets = ["Linear", "Conv1d", "Conv2d", "Conv3d", "QATLinear", "QATConv2d", "QATConv3d", "Conv1D"]
+            prunable_targets = get_prunable_targets()
             for name, module in match_named_modules(layer, prunable_targets).items():
                 name = f"{layer_name}.{name}"
 
@@ -208,7 +211,7 @@ def _infer_owl_layer_sparsity(
 
         groups = {}
         for name, layer in layers.items():
-            prunable_targets = ["Linear", "Conv1d", "Conv2d", "Conv3d", "QATLinear", "QATConv2d", "QATConv3d", "Conv1D"]
+            prunable_targets = get_prunable_targets()
             prunable_layers = match_named_modules(layer, prunable_targets)
             z = [
                 m.weight.abs() * activations[f"{name}.{n}"].unsqueeze(0)

src/llmcompressor/modifiers/smoothquant/base.py

@@ -206,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 =match_modules_set(
+                        _, balance_layer = match_modules_set(
                             balance_suffix, layer_name, model
                         )
                         if balance_layer:

src/llmcompressor/transformers/compression/helpers.py

@@ -76,8 +76,8 @@ 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 = match_named_modules(model, linear=True)
         offloaded_params = get_state_dict_offloaded_model(model)
+        linear_modules = match_named_modules(model, ["Linear"])
 
         linear_modules_with_sparsity_structure = [
             tensor_follows_mask_structure(offloaded_params[f"{name}.weight"])