change QuantizationModifier functions

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

Author: kylesayrs

src/llmcompressor/modifiers/modifier.py

@@ -89,7 +89,7 @@ def initialize(self, state: State, **kwargs):
 
         self.initialized_ = self.on_initialize(state=state, **kwargs)
 
-        # trigger starts
+        # trigger start
         fake_start_event = Event(type_=EventType.BATCH_START, global_step=0)
         if self.should_start(fake_start_event):
             self.on_start(state, fake_start_event, **kwargs)

src/llmcompressor/modifiers/quantization/gptq/base.py

@@ -3,6 +3,7 @@
 from typing import Dict, List, Optional, Tuple, Union
 
 import torch
+from compressed_tensors.quantization import disable_quantization
 from compressed_tensors.utils import (
     align_module_device,
     get_execution_device,
@@ -14,10 +15,6 @@
 
 from llmcompressor.core import State
 from llmcompressor.modifiers import Modifier
-from llmcompressor.modifiers.quantization.calibration import (
-    apply_calibration_status,
-    freeze_module_quantization,
-)
 from llmcompressor.modifiers.quantization.gptq.gptq_quantize import (
     accumulate_hessian,
     make_empty_hessian,
@@ -139,8 +136,13 @@ def on_initialize(self, state: State, **kwargs) -> bool:
         """
         # apply config to model and prepare calibration hooks
         if QuantizationMixin.has_config(self):
-            QuantizationMixin.attach_scheme_and_observers(self, state.model)
-            QuantizationMixin.register_calibration_hooks(self, state.model)
+            QuantizationMixin.initialize_quantization(self, state.model)
+
+        # assume quantization has been initialized by this modifier or one before it
+        QuantizationMixin.start_calibration(self, state.model)
+        # Unlike qmod, do not quantize as we calibrate
+        # This choice does not seem to have a meaningful impact on accuracy
+        state.model.apply(disable_quantization)
 
         # prepare module names
         self._module_names = {m: name for name, m in state.model.named_modules()}
@@ -162,9 +164,6 @@ def on_initialize(self, state: State, **kwargs) -> bool:
                 "modifier or a modifier preceding it"
             )
 
-        # prepare for calibration
-        state.model.apply(apply_calibration_status)
-
         # infer sequential targets
         if self.sequential_targets is None:
             self.sequential_targets = get_no_split_params(state.model)
@@ -233,8 +232,8 @@ def on_finalize(self, state: State, **kwargs) -> bool:
         self._hessians = dict()
         self._num_samples = dict()
 
-        state.model.apply(freeze_module_quantization)  # remove observers
-        self.remove_hooks()  # remove hooks
+        QuantizationMixin.end_calibration(self, state.model)
+        self.remove_hooks()  # remove gptq hooks
 
         return True
 

src/llmcompressor/modifiers/quantization/quantization/base.py

@@ -1,15 +1,10 @@
 import torch
 import tqdm
-from compressed_tensors.quantization import disable_quantization, enable_quantization
 from loguru import logger
 
 from llmcompressor.core import Event, State
 from llmcompressor.modifiers import Modifier
-from llmcompressor.modifiers.quantization.calibration import (
-    apply_calibration_status,
-    freeze_module_quantization,
-    update_weight_zp_scale,
-)
+from llmcompressor.modifiers.quantization.calibration import update_weight_zp_scale
 from llmcompressor.modifiers.quantization.quantization.mixin import QuantizationMixin
 from llmcompressor.modifiers.utils.pytorch_helpers import run_calibration_forward
 from llmcompressor.utils.helpers import calibration_forward_context
@@ -64,8 +59,7 @@ def on_initialize(self, state: State, **kwargs) -> bool:
                 "QuantizationModifier requires that quantization fields to be specified"
             )
 
-        QuantizationMixin.attach_scheme_and_observers(self, state.model)
-        state.model.apply(disable_quantization)  # disable quantization until start
+        QuantizationMixin.initialize_quantization(self, state.model)
 
         # FUTURE: modify oneshot lifecycle to trigger on_start for on initialize
         if self.calculate_start() == -1:  # one shot
@@ -77,9 +71,7 @@ def on_start(self, state: State):
         """
         Begin calibrating activations and weights. Calibrate weights only once on start
         """
-        QuantizationMixin.register_calibration_hooks(self, state.model)
-        state.model.apply(apply_calibration_status)
-        state.model.apply(enable_quantization)
+        QuantizationMixin.start_calibration(self, state.model)
 
         modules = list(state.model.modules())
         for module in tqdm.tqdm(modules, desc="Calibrating weights"):
@@ -93,8 +85,9 @@ def on_end(self, state: State, event: Event, **kwargs):
         """
         Finish calibrating by removing observers and calibration hooks
         """
-        state.model.apply(freeze_module_quantization)  # remove observers
-        self.remove_hooks()  # remove hooks
+        QuantizationMixin.end_calibration(
+            self, state.model
+        )  # keep quantization enabled
 
     def on_finalize(self, state: State, **kwargs) -> bool:
         # TODO: modify lifecycle so modifiers end on finalize

src/llmcompressor/modifiers/quantization/quantization/mixin.py

@@ -7,17 +7,22 @@
     QuantizationScheme,
     QuantizationStatus,
     apply_quantization_config,
+    disable_quantization,
+    enable_quantization,
     is_attention_module,
     is_preset_scheme,
     preset_name_to_scheme,
 )
-from pydantic import Field, field_validator
+from pydantic import Field, PrivateAttr, field_validator
+from torch.utils.hooks import RemovableHandle
 
 from llmcompressor.modifiers.quantization.calibration import (
+    apply_calibration_status,
     calibrate_input_hook,
     calibrate_kv_cache_input_hook,
     calibrate_kv_cache_output_hook,
     calibrate_output_hook,
+    freeze_module_quantization,
     initialize_observer,
     initialize_quantized_kv_cache,
     reset_quantization_status,
@@ -33,18 +38,18 @@ class QuantizationMixin(HooksMixin):
     calibration hooks, and compression wrappers to modifiers
 
     Lifecycle:
-        - QuantizationMixin.attach_scheme_and_observers(model)
-            - Wraps model forward and attaches quantization scheme and observers
-        - QuantizationMixin.register_calibration_hooks(model)
-            - Registers calibration hooks which utilize observers to calibrate qparams
-        - model.apply(apply_calibration_status)
-        - [ Calibrate model ]
-        - model.apply(freeze_module_quantization)
-            - Remove observers
-        - self.remove_hooks()
+        - on_initialize: QuantizationMixin.initialize_quantization
+            - Attach schemes to modules
+            - Attach observers to modules
+            - Disable quantization until calibration starts/finishes
+        - on_start: QuantizationMixin.start_calibration
+            - Attach calibration hooks
+            - Apply calibration status
+            - Enable quantization during calibration
+        - on_end: QuantizationMixin.end_calibration
             - Remove calibration hooks
-
-        Scheme is left attached to modules after PTQ finishes
+            - Apply freeze status
+            - Keep quantization enabled for future steps
 
     :param config_groups: dictionary specifying quantization schemes to apply to target
         modules. Modules not matching a scheme target will NOT be quantized.
@@ -76,6 +81,8 @@ class QuantizationMixin(HooksMixin):
     scheme: Optional[Union[str, Dict[str, Any]]] = None
     kv_cache_scheme: Optional[QuantizationArgs] = None
 
+    _calibration_hooks: List[RemovableHandle] = PrivateAttr(default_factory=list)
+
     @field_validator("targets", mode="before")
     def validate_targets(cls, value: Union[str, List[str]]) -> List[str]:
         if isinstance(value, str):
@@ -102,25 +109,49 @@ def validate_scheme(
 
         return value
 
-    def attach_scheme_and_observers(self, model: torch.nn.Module):
+    def initialize_quantization(self, model: torch.nn.Module):
         """
-        Apply this modifier as a quantization config to the model. Attach observers
-        according to the schemes attached to each module
+        Attach quantization schemes and observers to modules in the model according to
+        the quantization config specified on this modifier
+
+        :param model: model to attach schemes and observers to
         """
         reset_quantization_status(model)  # reset any previously applied qconfigs
 
+        # apply scheme and status to model
         config = self.resolve_quantization_config()
         apply_quantization_config(model, config)
 
+        # apply observers, disable quantization until calibration
         model.apply(self._initialize_observers)
+        model.apply(disable_quantization)
+
+    def start_calibration(self, model: torch.nn.Module):
+        """
+        Register activation calibration hooks (including kv_cache quantization) and
+        enable quantization as we calibrate
+
+        :param model: model to prepare for calibration
+        """
+        self._calibration_hooks = self._initialize_hooks(model)
+        model.apply(apply_calibration_status)
+        model.apply(enable_quantization)  # quantize at the same time as calibrate
 
-    def register_calibration_hooks(self, model: torch.nn.Module):
+    def end_calibration(self, model: torch.nn.Module):
         """
-        Register activation calibration hooks (including kv_cache quantization)
+        Remove calibration hooks and set the model status to frozen. Keep quantization
+        enabled for future operations
+
+        :param model: model to end calibration for
         """
-        model.apply(self._initialize_hooks)
+        self.remove_hooks(self._calibration_hooks)
+        model.apply(freeze_module_quantization)  # remove observers
+        model.apply(enable_quantization)  # keep quantization enabled
 
     def has_config(self) -> bool:
+        """
+        Determine if the user has specified a quantization config on this modifier
+        """
         return not (
             self.config_groups is None
             and self.targets == ["Linear"]
@@ -199,27 +230,44 @@ def _initialize_observers(self, module: torch.nn.Module):
         elif output:
             initialize_observer(module, base_name="output")
 
-    def _initialize_hooks(self, module: torch.nn.Module):
-        if not hasattr(module, "quantization_scheme"):
-            return
-
-        scheme: QuantizationScheme = module.quantization_scheme
-        input = scheme.input_activations and not scheme.input_activations.dynamic
-        output = scheme.output_activations and not scheme.output_activations.dynamic
-        is_attention = is_attention_module(module)
-
-        # input activations
-        if input:
-            self.register_hook(module, calibrate_input_hook, "forward_pre")
+    def _initialize_hooks(self, model: torch.nn.Module) -> List[RemovableHandle]:
+        hooks = []
+        for module in model.modules():
+            if not hasattr(module, "quantization_scheme"):
+                continue
+
+            scheme: QuantizationScheme = module.quantization_scheme
+            input = scheme.input_activations and not scheme.input_activations.dynamic
+            output = scheme.output_activations and not scheme.output_activations.dynamic
+            is_attention = is_attention_module(module)
+
+            # input activations
+            if input:
+                hooks.append(
+                    self.register_hook(module, calibrate_input_hook, "forward_pre")
+                )
+
+            # kv_cache activations. Within `apply_quantization_config`, the config is
+            # modified to use attention output quantization if a kv_cache_scheme exists
+            if is_attention and output:
+                hooks.append(
+                    self.register_hook(
+                        module,
+                        calibrate_kv_cache_input_hook,
+                        "forward_pre",
+                        with_kwargs=True,
+                    )
+                )
+                hooks.append(
+                    self.register_hook(
+                        module, calibrate_kv_cache_output_hook, "forward"
+                    )
+                )
 
-        # kv_cache activations. Within `apply_quantization_config`, the config is
-        # modified to use attention output quantization if a kv_cache_scheme exists
-        if is_attention and output:
-            self.register_hook(
-                module, calibrate_kv_cache_input_hook, "forward_pre", with_kwargs=True
-            )
-            self.register_hook(module, calibrate_kv_cache_output_hook, "forward")
+            # output activations
+            elif output:
+                hooks.append(
+                    self.register_hook(module, calibrate_output_hook, "forward")
+                )
 
-        # output activations
-        elif output:
-            self.register_hook(module, calibrate_output_hook, "forward")
+        return hooks