scoped quant status/config

Signed-off-by: Brian Dellabetta <[email protected]>

Author: brian-dellabetta

src/llmcompressor/modifiers/awq/base.py

@@ -265,8 +265,10 @@ def on_end(self, state: State, event: Event, **kwargs):
 
         self.ended_ = True
 
-        modules = list(state.model.modules())
-        for module in tqdm(modules, desc="Calibrating weights"):
+        for _, module in tqdm(
+            match_named_modules(state.model, self.targets, self.ignore),
+            desc="Calibrating weights",
+        ):
             update_weight_zp_scale(module)
 
         QuantizationMixin.end_calibration(self, state.model)

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

@@ -1,7 +1,6 @@
 import tqdm
-from compressed_tensors.utils import (
-    match_named_modules,
-)
+from compressed_tensors.utils import match_named_modules
+
 from llmcompressor.core import Event, EventType, State
 from llmcompressor.modifiers import Modifier
 from llmcompressor.modifiers.quantization.calibration import (
@@ -77,10 +76,10 @@ def on_start(self, state: State, event: Event, **kwargs):
         # TODO: this step can be combined with update_weight_zp_scale
         # once update_fused_layer_weight_global_scales is removed
         # and not required by vLLM
-        for name, module in tqdm.tqdm(named_modules):
+        for _, module in tqdm.tqdm(named_modules):
             update_weight_global_scale(module)
 
-        for name, module in tqdm.tqdm(named_modules, desc="Calibrating weights"):
+        for _, module in tqdm.tqdm(named_modules, desc="Calibrating weights"):
             update_fused_layer_weight_global_scales(module)
             update_weight_zp_scale(module)
 

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

@@ -14,6 +14,7 @@
     is_preset_scheme,
     preset_name_to_scheme,
 )
+from compressed_tensors.utils import match_named_modules
 from pydantic import Field, PrivateAttr, field_validator
 from torch.utils.hooks import RemovableHandle
 
@@ -121,12 +122,15 @@ def initialize_quantization(self, model: torch.nn.Module):
 
         :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()
+
+        for _, module in match_named_modules(model, self.targets, self.ignore):
+            reset_quantization_status(module)  # reset any previously applied qconfigs
+
         apply_quantization_config(model, config)
 
+        # TODO should we disable for entire model or just matching modules?
         # disable quantization until calibration
         model.apply(disable_quantization)
 
@@ -138,8 +142,11 @@ def start_calibration(self, model: torch.nn.Module):
         :param model: model to prepare for calibration
         """
         self._calibration_hooks = self._initialize_hooks(model)
-        model.apply(self._initialize_observers)
-        model.apply(apply_calibration_status)
+        for _, module in match_named_modules(model, self.targets, self.ignore):
+            self._initialize_observers(module)
+            apply_calibration_status(module)
+
+        # TODO should we disable for entire model or just matching modules?
         model.apply(enable_quantization)  # quantize at the same time as calibrate
 
     def end_calibration(self, model: torch.nn.Module):
@@ -150,7 +157,9 @@ def end_calibration(self, model: torch.nn.Module):
         :param model: model to end calibration for
         """
         self.remove_hooks(self._calibration_hooks)
-        model.apply(freeze_module_quantization)  # remove observers
+        for _, module in match_named_modules(model, self.targets, self.ignore):
+            freeze_module_quantization(module)  # remove observers
+
         model.apply(enable_quantization)  # keep quantization enabled
 
     def has_config(self) -> bool:
@@ -240,7 +249,7 @@ def _initialize_observers(self, module: torch.nn.Module):
 
     def _initialize_hooks(self, model: torch.nn.Module) -> Set[RemovableHandle]:
         hooks = set()
-        for module in model.modules():
+        for _, module in match_named_modules(model, self.targets, self.ignore):
             if not hasattr(module, "quantization_scheme"):
                 continue