Extend disable_hooks to keep subsets

[kylesayrs]

Purpose

• Allow subsets of handles to remain active, as is needed in the case of wanda

def _get_activations(self, model, dataloader, nsamples=128):
  def save_acts(module, input, name):
      ...
  
  hooks = set(
      self.register_hook(mod, partial(save_acts, name=name), "forward_pre")
      for name, mod in self.model.named_modules()
      if isinstance(mod, torch.nn.Linear) and "lm_head" not in name
  )
  # in the future, if the user puts wanda after another modifier,
  # initialize will run after other modifiers have added hooks

  # want to disable hooks from other modifiers, but keep the ones we just added
  with HooksMixin.disable_hooks(keep=hooks):
      run_basic(model, dataloader)
  self.remove_hooks(hooks)

Prerequisites

• Extend remove_hooks to remove subsets #1021

Postrequisites

• Layer compressor deprecation

Changes

• Add a _HOOKS_KEEP_ENABLED class variable

Tests

• Added test_disable_hooks_keep in tests

[github-actions]

👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review.

[rahul-tuli]

The new changes LGTM! Great work @kylesayrs

The base branch was changed.