formatting

Summary

Signed-off-by: HDCharles <[email protected]>

Author: HDCharles

src/llmcompressor/modifiers/awq/base.py

@@ -362,12 +362,11 @@ def _set_resolved_mappings(self, model: Module) -> None:
                     )
 
                     continue
-                else:
-                    # for multiple balance layers, find lowest common parent
-                    ancestor_name = get_lowest_common_ancestor_name(balance_names)
-                    ancestor_name, ancestor = get_lowest_non_module_list_ancestor(
-                        ancestor_name, model
-                    )
+
+                ancestor_name = get_lowest_common_ancestor_name(balance_names)
+                ancestor_name, ancestor = get_lowest_non_module_list_ancestor(
+                    ancestor_name, model
+                )
 
                 resolved_mappings.append(
                     ResolvedMapping(

src/llmcompressor/utils/pytorch/module.py

@@ -378,8 +378,7 @@ def get_module_to_name_dict(model: Module) -> dict[Module:str]:
         if module in module_to_name:
             logger.info(
                 f"Warning, {name} and {module_to_name[module]} both "
-                "share the same module the same module, "
-                "may have trouble resolving mappings."
+                "share the same module, which can result in unexpected behavior"
             )
         module_to_name[module] = name
     return module_to_name

tests/llmcompressor/pytorch/modifiers/logarithmic_equalization/test_pytorch.py

@@ -19,4 +19,4 @@ def test_log_equalization_mapping(state):
     mapping = modifier.resolved_mappings_[0]
     assert mapping.smooth_name == mappings[0][1]
     assert isinstance(mapping.smooth_layer, Linear)
-    assert isinstance(mapping.balance_layers[0], Linear)
+    assert isinstance(mapping.balance_layers[0], Linear)

tests/llmcompressor/pytorch/modifiers/smoothquant/test_pytorch.py

@@ -4,12 +4,6 @@
 from llmcompressor.modifiers.smoothquant import SmoothQuantModifier
 
 
-import pytest
-from torch.nn import Linear
-
-from llmcompressor.modifiers.smoothquant import SmoothQuantModifier
-
-
 @pytest.mark.unit
 def test_smooth_quant_mapping(state):
     mappings = [(["seq.fc1"], "seq.fc2")]
@@ -23,4 +17,4 @@ def test_smooth_quant_mapping(state):
     mapping = modifier.resolved_mappings_[0]
     assert mapping.smooth_name == mappings[0][1]
     assert isinstance(mapping.smooth_layer, Linear)
-    assert isinstance(mapping.balance_layers[0], Linear)
+    assert isinstance(mapping.balance_layers[0], Linear)