Address review feedback for vllm-project#1687: use Module.get_submodule; fix match_named_modules arg order; clean imports

Author: EtelisIBM

src/llmcompressor/modifiers/awq/base.py

@@ -14,7 +14,6 @@
 from loguru import logger
 from pydantic import ConfigDict, PrivateAttr, model_validator
 from torch.nn import Module
-from operator import attrgetter
 from tqdm import tqdm
 
 from llmcompressor.core import Event, EventType, State
@@ -324,14 +323,16 @@ 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) if smooth_parent_name else model
+                smooth_parent = (
+                    model if smooth_parent_name == "" else model.get_submodule(smooth_parent_name)
+                )
 
                 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 match_named_modules(
-                        balance_regex,
                         smooth_parent,
+                        balance_regex,
                         exclude_internal_modules=True,
                     ).items():
                         balance_name = f"{smooth_parent_name}.{balance_suffix}"
@@ -766,7 +767,7 @@ def get_lowest_common_parent(names: List[str], module: Module) -> Tuple[str, Mod
     while True:
         if parent_name == "":
             return "", module
-        parent = attrgetter(parent_name)(module)
+        parent = module.get_submodule(parent_name)
         if not isinstance(parent, torch.nn.ModuleList):
             return parent_name, parent
         parent_name = ".".join(parent_name.split(".")[:-1])

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

@@ -61,8 +61,8 @@ def on_initialize(self, state: State, **kwargs) -> bool:
             else:
                 model_target, teacher_target = target, target
 
-            model_layers = match_named_modules(model_target, state.model)
-            teacher_layers = match_named_modules(teacher_target, state.teacher_model)
+            model_layers = match_named_modules(state.model, model_target)
+            teacher_layers = match_named_modules(state.teacher_model, teacher_target)
 
             if len(model_layers) < 1:
                 raise ValueError(f"no model layers found for target {target}")