Implement QuantizationMixin (#1351)

## Purpose ## 
* Abstract functionality which allows modifiers to act as quantization
configs into a mixin called `QuantizationMixin`
* This gives #1279 an interface to properly infer which pipeline to use
based on the recipe (if a recipe contains modifiers requires
calibration, then use the "basic" or "sequential" pipelines)
* This enables future modifiers to act as quantization modifiers (in the
same way that GPTQ does now)
* Related to #1354 where previous logic would attempt to add a
QuantizedKVCache for dynamic kv_quant

## Changes ##
* Implement `QuantizationMixin` which implements five public methods
  * Lifecycle methods
* `initialize_quantization` is used to apply a config and attach
observers to a model
* quantization is disabled so that modules aren't quantized before
they're calibrated
* `start_calibration` is used to initialize calibration hooks and status
* quantization is enabled, since we currently quantize as we calibrate,
although this decision is somewhat arbitrary
* `end_calibration` is used to remove calibration hooks and apply the
frozen status
* quantization remains enabled, since we want future forward passes to
simulate quantization
  * Recipe-related methods
* `has_config` returns true if a config was specified, used for checking
against duplicate configs in the recipe
* `resolve_quantization_config` returns the quantization config
specified by the modifier fields
* `QuantizationModifier` inherits from `QuantizationMixin`
* `GPTQModifier` inherits from `QuantizationMixin`
* Unlike QMod, GPTQ disables quantization during calibration. As noted
before, this is a somewhat arbitrary choice but one which matches the
current implementation

* Calibration utils
* Replace `set_unset_kv_cache` with `initialize_quantized_kv_cache` and
`freeze_module_quantization`
    * Treat the `QuantizedKVCache` as analogous to another observer
  * Pull setting the calibration status out of`update_weight_zp_scale`
* This better matches the lifecycle detailed in `QuantizationMixin`
description
* Implement `reset_quantization_status` which is used to remove any
existing quantization configs before the current config is applied by
`initialize_quantization`

## Remove Support ##
* Removing support for recipe with multiple quantization modifiers
active at the same time (a check for this will be added by #1279)
* Remove `num_calibration_steps`, `quantize`,
`disable_quantization_observer_epoch` and `min_tokens_per_module`
* `num_calibration_steps` is already controlled by
https://github.com/vllm-project/llm-compressor/blob/42b62f5283d0234b26623fe1f1bf02a77c6e4019/src/llmcompressor/datasets/utils.py#L106
* `quantize` was implemented as a workaround for GPTQ's modifier
builder. Similar functionality may be require to support SpinQuant +
GPTQ, but such functionality should exist at a higher level
* `disable_quantization_observer_epoch` seems to implement functionality
where a model's observers are removed but quantization remains active.
This functionality is maintained by setting an "end" epoch for qmod
* `min_tokens_per_module` requires that the modifier have references to
the calibration dataset, which is disallowed by #1279. This information
is already printed in GPTQ's logs. If research still wants this tool
specifically for `QuantizationModifier`, then it can be reimplemented to
avoid using references to the calibration dataset
  
## Testing ##
* Updated tests to reflect new mixin
* Ran a set of GPTQ and QuantizationModifier examples to completion
* CI tests pass

---------

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

Author: kylesayrs

src/llmcompressor/modifiers/quantization/calibration.py

@@ -3,8 +3,8 @@
 import torch
 from compressed_tensors.quantization import (
     KVCacheScaleType,
+    QuantizationScheme,
     QuantizationStatus,
-    is_attention_module,
 )
 from compressed_tensors.quantization.lifecycle.forward import forward_quantize
 from compressed_tensors.quantization.utils import is_kv_cache_quant_scheme
@@ -14,6 +14,7 @@
 
 from llmcompressor.modifiers.quantization.cache import QuantizedKVParameterCache
 from llmcompressor.observers import Observer
+from llmcompressor.utils.helpers import getattr_chain
 
 __all__ = [
     "initialize_observer",
@@ -22,9 +23,10 @@
     "calibrate_output_hook",
     "calibrate_kv_cache_input_hook",
     "calibrate_kv_cache_output_hook",
-    "set_unset_kv_cache",
+    "initialize_quantized_kv_cache",
     "freeze_module_quantization",
     "apply_calibration_status",
+    "reset_quantization_status",
 ]
 
 
@@ -49,10 +51,6 @@ def initialize_observer(
         # no quantization scheme nothing to do
         return
 
-    # observers have a different lifecycle for kv_cache
-    if is_attention_module(module):
-        return
-
     quantization_args = getattr(quantization_scheme, arg_name, None)
     # dont need observers for dynamic
     if quantization_args is not None and not quantization_args.dynamic:
@@ -102,25 +100,15 @@ def update_weight_zp_scale(module: Module):
     :param quantize_weights_upfront: whether to automatically
        run weight quantization at the start of calibration
     """
-    if not getattr(module, "quantization_scheme", None):
-        # no quantization scheme nothing to do
+    if getattr_chain(module, "quantization_scheme.weights", None) is None:
         return
 
-    status = getattr(module, "quantization_status", None)
-    if not status:
-        # not set to initialize; no scales/zp to update
-        return
-    if status != QuantizationStatus.INITIALIZED:
+    if getattr(module, "quantization_status", None) != QuantizationStatus.CALIBRATION:
         logger.warning(
-            f"Attempting set module with status {status} to calibration mode. "
-            f"but status is not {QuantizationStatus.INITIALIZED} - you may "
-            "be calibrating an uninitialized module which may fail or attempting "
-            "to re-calibrate a frozen module"
+            "Attempting to calibrate weights of a module not in calibration mode"
         )
 
-    if module.quantization_scheme.weights is not None:
-        # set weight scale and zero_point up front, calibration data doesn't affect it
-        call_observer(module=module, base_name="weight")
+    call_observer(module=module, base_name="weight")
 
 
 def calibrate_activations(module: Module, value: torch.Tensor, base_name: str):
@@ -200,21 +188,26 @@ def calibrate_kv_cache_output_hook(module: Module, _args: Any, _output: torch.Te
     update_parameter_data(module, v_scale, KVCacheScaleType.VALUE.value)
 
 
-def set_unset_kv_cache(module: Module):
+def initialize_quantized_kv_cache(module: Module):
     """
-    Set or unset singleton QuantizedKVParameterCache for each
-    attn module when running kv_cache quantization.
+    Initialize a quantized kv_cache on a module (analogous to initializing an observer)
+    When a config specifying kv_cache quantization is applied to a model, the kv_cache
+    args are redefined as the output_activations targeting attention modules.
+
+    This function should be called on attention modules with output_activations
     """
-    if not hasattr(module, "quantization_scheme"):
+    scheme: Optional[QuantizationScheme] = getattr(module, "quantization_scheme", None)
+    existing_kv_cache = getattr(module, "kv_cache", None)
+
+    if (
+        scheme is None
+        or not is_kv_cache_quant_scheme(scheme)
+        or isinstance(existing_kv_cache, QuantizedKVParameterCache)
+    ):
         return
 
-    if is_kv_cache_quant_scheme(module.quantization_scheme):
-        output_args = module.quantization_scheme.output_activations
-        kv_cache = QuantizedKVParameterCache(output_args)
-        if hasattr(module, "kv_cache"):
-            delattr(module, "kv_cache")
-        else:
-            setattr(module, "kv_cache", kv_cache)
+    quantized_kv_cache = QuantizedKVParameterCache(scheme.output_activations)
+    setattr(module, "kv_cache", quantized_kv_cache)
 
 
 def apply_calibration_status(module: Module):
@@ -242,9 +235,21 @@ def freeze_module_quantization(module: Module):
         # nothing to do, already frozen
         return
 
+    # remove observers
     for name in ("input", "weight", "output"):
         obs_name = f"{name}_observer"
         if hasattr(module, obs_name):
             delattr(module, obs_name)
 
+    # remove quantized kv_cache
+    kv_cache = getattr(module, "kv_cache", None)
+    if isinstance(kv_cache, QuantizedKVParameterCache):
+        delattr(module, "kv_cache")
+
     module.quantization_status = QuantizationStatus.FROZEN
+
+
+def reset_quantization_status(model: Module):
+    for module in model.modules():
+        if hasattr(module, "quantization_status"):
+            delattr(module, "quantization_status")