Extend remove_hooks to remove subsets (#1021)

## Purpose ##
* Allow subsets of hooks to be removed
* Not strictly needed but helps promote code clarity in the case of
wanda which adds and removes subsets of hooks at different times.

## Postrequisites ##
* #1023
* Layer compressor deprecation

## Changes ##
* Change the datatype of `_hooks` from `List` to `Set`
* Add `handles` argument to `HooksMixin.remove_hooks`

## Testing ##
* Added `test_remove_hooks_parameterized` test

---------

Signed-off-by: Kyle Sayers <[email protected]>

Author: kylesayrs

src/llmcompressor/modifiers/utils/hooks.py

@@ -1,6 +1,6 @@
 import contextlib
 from functools import wraps
-from typing import Any, Callable, ClassVar, List, Union
+from typing import Any, Callable, ClassVar, Optional, Set, Union
 
 import torch
 from loguru import logger
@@ -30,7 +30,7 @@ class HooksMixin(BaseModel):
     """
 
     _HOOKS_DISABLED: ClassVar[bool] = False  # attached to global HooksMixin
-    _hooks: List[RemovableHandle] = []  # attached to local subclasses
+    _hooks: Set[RemovableHandle] = set()  # attached to local subclasses
 
     @classmethod
     @contextlib.contextmanager
@@ -70,14 +70,22 @@ def wrapped_hook(*args, **kwargs):
 
         register_function = getattr(target, f"register_{hook_type}_hook")
         handle = register_function(wrapped_hook, **kwargs)
-        self._hooks.append(handle)
+        self._hooks.add(handle)
         logger.debug(f"{self} added {handle}")
 
         return handle
 
-    def remove_hooks(self):
-        """Remove all hooks belonging to a modifier"""
-        for hook in self._hooks:
+    def remove_hooks(self, handles: Optional[Set[RemovableHandle]] = None):
+        """
+        Removes hooks registered by this modifier
+
+        :param handles: optional list of handles to remove, defaults to all hooks
+            registerd by this modifier
+        """
+        if handles is None:
+            handles = self._hooks
+
+        for hook in handles:
             hook.remove()
 
-        self._hooks = []
+        self._hooks -= handles

tests/llmcompressor/modifiers/utils/test_hooks.py

@@ -64,6 +64,27 @@ def test_remove_hooks():
     assert mod_a.hook_called and not mod_b.hook_called
 
 
+def test_remove_hooks_parameterized():
+    model = DummyModel()
+
+    mod_a = ModA()
+    mod_a_pre_hook = mod_a.register_hook(model.linear1, mod_a.hook, "forward_pre")
+    mod_a_post_hook = mod_a.register_hook(model.linear1, mod_a.hook, "forward")
+
+    mod_b = ModB()
+    mod_b_pre_hook = mod_b.register_hook(model.linear2, mod_b.hook, "forward_pre")
+    mod_b_post_hook = mod_b.register_hook(model.linear2, mod_b.hook, "forward")
+
+    mod_a.remove_hooks(set([mod_a_post_hook]))
+    mod_b.remove_hooks(set([mod_b_pre_hook]))
+
+    assert len(mod_a._hooks) == 1 and next(iter(mod_a._hooks)) == mod_a_pre_hook
+    assert len(mod_b._hooks) == 1 and next(iter(mod_b._hooks)) == mod_b_post_hook
+
+    model(model.dummy_inputs)
+    assert mod_a.hook_called and mod_b.hook_called
+
+
 def test_disable_hooks():
     model = DummyModel()