format

Summary

Signed-off-by: HDCharles <charlesdavidhernandez@gmail.com>

Author: HDCharles

src/llmcompressor/modifiers/awq/base.py

@@ -32,7 +32,6 @@
 from llmcompressor.utils.fsdp.helpers import get_fsdp_parent
 from llmcompressor.utils.helpers import calibration_forward_context
 from llmcompressor.utils.pytorch.module import (
-    get_layer_by_name,
     get_module_to_name_dict,
 )
 
@@ -365,7 +364,10 @@ def _set_resolved_mappings(self, model: Module) -> None:
 
                 ancestor_name = get_lowest_common_ancestor_name(balance_names)
                 # no ModuleList ancestors
-                while not isinstance((ancestor := model.get_submodule(ancestor_name)), torch.nn.ModuleList):
+                while not isinstance(
+                    (ancestor := model.get_submodule(ancestor_name)),
+                    torch.nn.ModuleList,
+                ):
                     ancestor_name = ancestor_name.rsplit(".", 1)[0]
 
                 resolved_mappings.append(
@@ -798,28 +800,3 @@ def _accumulate_mean(
     new_count = prev_count + num_added
 
     return (prev_sum + sum_added) / new_count, new_count
-
-
-def get_lowest_non_module_list_ancestor(name, module: Module) -> tuple[str, Module]:
-    """
-    Given a name and a model, finds lowest ancestor of
-    named module that's not a ModuleList
-    i.e. module_list.module_dict.module_list -> module_list.module_dict
-    i.e. module_list.module_dict -> module_list.module_dict
-    (self is an ancestor of self)
-
-    NOTE: This is needed because ModuleLists don't play
-    nicely with hooks because their forward method is never directly
-    called for MoE models. See Qwen3MoeSparseMoeBlock for example, experts
-    are selected based on router output and their forward method is called.
-    https://github.com/huggingface/transformers/blob/v4.52.4/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py#L233
-
-    Returns name of module and pointer to module
-    """
-    while True:
-        if name == "":
-            return "", module
-        current_module = get_layer_by_name(name, module)
-        if not isinstance(current_module, torch.nn.ModuleList):
-            return name, current_module
-        name = ".".join(name.split(".")[:-1])

tests/llmcompressor/modifiers/awq/test_base.py

@@ -5,7 +5,6 @@
 from torch.nn import Linear
 
 from llmcompressor.modifiers.awq import AWQMapping, AWQModifier
-from llmcompressor.modifiers.awq.base import get_lowest_non_module_list_ancestor
 from llmcompressor.modifiers.factory import ModifierFactory
 
 
@@ -205,39 +204,6 @@ def test_validate():
         AWQModifier(scheme="W4A16", duo_scaling="x")
 
 
-@pytest.mark.unit
-def test_get_lowest_non_module_list_ancestor():
-    model = torch.nn.ModuleDict(
-        {
-            "experts": torch.nn.ModuleList(
-                [
-                    torch.nn.ModuleDict(
-                        {
-                            "gate_proj": Linear(4, 2),
-                            "down_proj": Linear(2, 4),
-                        }
-                    )
-                    for _ in range(10)
-                ]
-            )
-        }
-    )
-
-    ancestor_name, ancestor = get_lowest_non_module_list_ancestor("", model)
-    assert ancestor_name == "" and ancestor == model
-
-    ancestor_name, ancestor = get_lowest_non_module_list_ancestor("experts", model)
-    assert ancestor_name == "" and ancestor == model
-
-    ancestor_name, ancestor = get_lowest_non_module_list_ancestor(
-        "experts.1.gate_proj", model
-    )
-    assert (
-        ancestor_name == "experts.1.gate_proj"
-        and ancestor == model["experts"][1]["gate_proj"]
-    )
-
-
 @pytest.mark.unit
 def test_moe_multiple_balance_layers():
     """Test AWQ mapping with multiple balance layers in MoE architecture"""