Implement HooksMixin

[kylesayrs]

Purpose

• Precursor to VLM Support via GPTQ Hooks and Data Pipelines #914
• Create a shared API for adding hooks to modules
• Allow code which handles data pipelines to selectively disable hooks for certain passes. This will be needed in cases with custom datapipelines (GPTQ/Wanda/SparseGPTQ) and when multiple modifiers are active at the same time.
  • This is needed for GPTQ-style sequential algorithms which require one pass with hooks in order to accumulate the hessians and compress, and then a second pass without hooks in order to compute compressed (weight-quantized) outputs
  • This is also a tool for research users to be able to control when hooks are enabled from within the data pipelines

for layer in model_layers:
    # accumulate hessians
    unquantized_outputs = layer(*args, **kwargs)

    # get sequential outputs
    with HooksMixin.disable_hooks():
        quantized_outputs = layer(*args, **kwargs)
    
    print(f"Mean error from quantization: {get_loss(unquantized_outputs, quantized_outputs)}")

Changes

• Implement HooksMixin
  • The _HOOKS_DISABLED attribute is a global variable attached to the class which is used to disable hooks globally
  • The _hooks attribute is a local variable attached to each modifier which lists all of the hooks created by that modifier
• Integrate with QuantizationModifier, refactor calibration functions to reference the same function rather than generating hook functions
• Integrate with SmoothQuantModifier
• Integrate with WandaPruningModifier and SparseGPTModifier
• Integrate with MagnitudePruningModifier and ConstantPruningModifier via LayerParamMasking
• Purposefully did not integrate with LayerCompressor since this will be handled by future data pipelines and doing so would all the BaseModel inheritance to the LayerCompressor class, which add unnecessary complexity to this PR

Testing

• Added tests in tests/llmcompressor/modifiers/utils/test_hooks.py

[github-actions]

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

[kylesayrs]

e2e tests
nightly: https://github.com/neuralmagic/llm-compressor-testing/actions/runs/11897900649 ✅

[dsikka]

We briefly looked at the implications of using hooks with FSDP - are we taking care of that already or through this PR?

[kylesayrs]

@dsikka I consider that to be out of scope for this PR. I consider FSDP to be unsupported as of now, although this PR makes it easier to support FSDP in the future.

Modifying a module's parameter requires being in special FSDP contexts.

@torch.no_grad()
def pre_hook(module, _args):
  # modifying both training and handle training states is required
  with model._use_training_state(TrainingState.IDLE, HandleTrainingState.IDLE):
    with FullyShardedDataParallel.summon_full_params(model):
      # modify module weight. Doing so outside of the contexts will raise a non-contiguous tensor error
      module.weight *= 0

We can bake these contexts into the HooksMixin.register_hook function, although there's implementation details associated with that I'd like to leave that for a separate task/PR.

[dsikka]

Overall looks good in cleaning up/unifying hooks

Current testing should test the changes with the QuantizationModifier- do we think this is the case for the other modifiers being tested?

The other thought I had was about a less common but potentially useful use case where a modifier may have hooks for different cases and may want to target turning off a specific subset as opposed to all of them - do we think the hooks mixin class can be extended easily to handle that?

[kylesayrs]

@dsikka

Current testing should test the changes with the QuantizationModifier- do we think this is the case for the other modifiers being tested?

I've tested with the e2e tests, although I can perform more rigorous testing if we think that's necessary.

The other thought I had was about a less common but potentially useful use case where a modifier may have hooks for different cases and may want to target turning off a specific subset as opposed to all of them - do we think the hooks mixin class can be extended easily to handle that?

Yes! There are good arguments to be made for enabling this kind of functionality within the GPTQ algorithm, and unifying hooks makes implementing this functionality much easier.

[dsikka]

I'd suggest checking out the nightly test cases and making sure we're not running any issues there. LGTM.

[dsikka]

oh ignore my nightly comment.