unit tests

Signed-off-by: Brian Dellabetta <bdellabe@redhat.com>

Author: brian-dellabetta

src/llmcompressor/modifiers/awq/base.py

@@ -749,17 +749,19 @@ def _accumulate_mean(
     return (prev_sum + sum_added) / new_count, new_count
 
 
-def get_lowest_common_parent(names: List[str], module: Module) -> str:
+def get_lowest_common_parent(names: List[str], module: Module) -> Tuple[str, Module]:
     """
     Given a list of names, returns the lowest-scope common parent,
     excluding parents of type ModuleList, which don't seem to play
     nicely with hooks.
-    Slight alteration from os.path.commonprefix
+    Returns name of parent and pointer to parent module
+
+    Implementation is a small alteration of os.path.commonprefix
     https://docs.python.org/3/library/os.path.html#os.path.commonprefix
     """
     s1 = min(names)
     s2 = max(names)
-    parent_name = s1
+    parent_name = ""
     for i, c in enumerate(s1):
         if c != s2[i]:
             parent_name = s1[:i].rstrip(".")

tests/llmcompressor/modifiers/awq/test_base.py

@@ -4,6 +4,7 @@
 from pydantic import ValidationError
 
 from llmcompressor.modifiers.awq import AWQMapping, AWQModifier
+from llmcompressor.modifiers.awq.base import get_lowest_common_parent
 from llmcompressor.modifiers.factory import ModifierFactory
 from tests.llmcompressor.modifiers.conf import setup_modifier_factory
 
@@ -172,3 +173,55 @@ def test_validate():
             ),
         }
     )
+
+
+@pytest.mark.unit
+def test_get_lowest_common_parent():
+    mlp = torch.nn.ModuleDict(
+        {
+            "experts": torch.nn.ModuleList(
+                [
+                    torch.nn.ModuleDict(
+                        {
+                            "gate_proj": torch.nn.Linear(4, 2),
+                            "down_proj": torch.nn.Linear(4, 2),
+                        }
+                    )
+                    for _ in range(10)
+                ]
+            )
+        }
+    )
+    self_attn = torch.nn.ModuleDict(
+        {
+            "q_proj": torch.nn.Linear(4, 2),
+            "k_proj": torch.nn.Linear(4, 2),
+            "v_proj": torch.nn.Linear(4, 2),
+            "o_proj": torch.nn.Linear(4, 4),
+        }
+    )
+    model = torch.nn.ModuleDict(
+        {
+            "decoder": torch.nn.ModuleDict(
+                {
+                    "self_attn": self_attn,
+                    "mlp": mlp,
+                }
+            )
+        }
+    )
+
+    parent_name, parent = get_lowest_common_parent(
+        ["decoder.mlp.experts.1.gate_proj", "decoder.mlp.experts.4.down_proj"], model
+    )
+    assert parent_name == "decoder.mlp" and parent == mlp
+
+    parent_name, parent = get_lowest_common_parent(
+        ["decoder.self_attn.q_proj", "decoder.self_attn.v_proj"], model
+    )
+    assert parent_name == "decoder.self_attn" and parent == self_attn
+
+    parent_name, parent = get_lowest_common_parent(
+        ["decoder.mlp.experts.1.gate_proj", "decoder.self_attn.v_proj"], model
+    )
+    assert parent_name == "decoder" and parent == model["decoder"]