get_layers_params refactor

Author: ved1beta

src/llmcompressor/modifiers/pruning/constant/base.py

@@ -8,8 +8,7 @@
     LayerParamMasking,
     param_mask_name,
 )
-from llmcompressor.utils.pytorch.module import get_layers_params
-
+from compressed_tensors import match_named_parameters
 __all__ = ["ConstantPruningModifier"]
 
 
@@ -29,7 +28,7 @@ def on_initialize(self, state: State, **kwargs) -> bool:
         if not state.model:
             return False
 
-        self.parameterized_layers_ = get_layers_params(self.targets, state.model)
+        self.parameterized_layers_ = match_named_parameters(self.targets, state.model)
 
         for layer_param_name, parameterized_layer in self.parameterized_layers_.items():
             self.add_mask(

src/llmcompressor/modifiers/pruning/magnitude/base.py

@@ -16,8 +16,7 @@
     PruningMaskCreatorArgs,
     PruningMaskFactory,
 )
-from llmcompressor.utils.pytorch.module import get_layers_params
-
+from compressed_tensors import match_named_parameters
 __all__ = ["MagnitudePruningModifier"]
 
 
@@ -73,7 +72,7 @@ def on_initialize(self, state: State, **kwargs) -> bool:
             self.mask_structure
         )
 
-        self.parameterized_layers_ = get_layers_params(state.model)
+        self.parameterized_layers_ = match_named_parameters(state.model)
 
         for layer_param_name, parameterized_layer in self.parameterized_layers_.items():
             self.add_mask(

src/llmcompressor/utils/pytorch/module.py

@@ -58,10 +58,8 @@
     "get_prunable_layers",
     "get_quantizable_layers",
     "qat_active",
-    "get_layers_params",
     "get_matching_layer",
     "get_no_split_params",
-    "get_layer_by_name",
 ]
 
 ALL_TARGET = "__ALL__"
@@ -292,22 +290,6 @@ def qat_active(module: Module) -> bool:
     return False
 
 
-def get_layers_params(
-    targets: Union[str, List[str]], module: Module
-) -> Dict[str, ModelParameterizedLayer]:
-    params = get_params(targets, module)
-    layers = get_layers(targets, module)
-
-    parameterized_layers = {}
-    for name, param in params.items():
-        param_layer = ModelParameterizedLayer(
-            layer_name=name, layer=layers[name], param_name=name, param=param
-        )
-        parameterized_layers[name] = param_layer
-
-    return parameterized_layers
-
-
 def get_matching_layer(
     target: str, name_to_match: str, module: Module
 ) -> Optional[Tuple[str, Module]]:

tests/llmcompressor/utils/pytorch/test_module.py

@@ -1,9 +1,6 @@
 import pytest
 import torch.nn as nn
 
-from llmcompressor.utils.pytorch import get_layer_by_name
-
-
 @pytest.fixture
 def example_nested_module() -> str:
     return nn.Sequential(
@@ -14,21 +11,3 @@ def example_nested_module() -> str:
     )
 
 
-@pytest.mark.unit
-def test_get_layer_by_name(example_nested_module):
-    # Test getting the parent of a nested layer
-    layer = get_layer_by_name("0", example_nested_module)
-    assert layer == example_nested_module[0]
-
-    layer = get_layer_by_name("1.1", example_nested_module)
-    assert layer == example_nested_module[1][1]
-
-    layer = get_layer_by_name("2.0", example_nested_module)
-    assert layer == example_nested_module[2][0]
-
-    layer = get_layer_by_name("2.1", example_nested_module)
-    assert layer == example_nested_module[2][1]
-
-    # Test getting the parent of a non-existent layer
-    with pytest.raises(AttributeError):
-        get_layer_by_name("non_existent_layer", example_nested_module)