match_named_modules, add observer on_start instead of on_initialize

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

Author: brian-dellabetta

src/llmcompressor/modifiers/awq/mappings.py

@@ -157,6 +157,7 @@ class AWQMapping:
     "Phi3ForCausalLM": _phi_mappings,
     "Phi3VForCausalLM": _phi_mappings,
     "Qwen2ForCausalLM": _default_mappings,
+    "Qwen2_5OmniThinkerForConditionalGeneration": _default_mappings,
     "Qwen2MoeForCausalLM": _moe_default_mappings,
     "Qwen3ForCausalLM": _default_mappings,
     "Qwen3MoeForCausalLM": _moe_default_mappings,

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

@@ -1,5 +1,7 @@
 import tqdm
-
+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 (
@@ -69,14 +71,16 @@ def on_start(self, state: State, event: Event, **kwargs):
         self.started_ = True
         QuantizationMixin.start_calibration(self, state.model)
 
-        modules = list(state.model.modules())
+        named_modules = list(
+            match_named_modules(state.model, self.targets, self.ignore)
+        )
         # 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 module in tqdm.tqdm(modules):
+        for name, module in tqdm.tqdm(named_modules):
             update_weight_global_scale(module)
 
-        for module in tqdm.tqdm(modules, desc="Calibrating weights"):
+        for name, 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

@@ -116,7 +116,7 @@ def validate_scheme(
 
     def initialize_quantization(self, model: torch.nn.Module):
         """
-        Attach quantization schemes and observers to modules in the model according to
+        Attach quantization schemes to modules in the model according to
         the quantization config specified on this modifier
 
         :param model: model to attach schemes and observers to
@@ -127,25 +127,25 @@ def initialize_quantization(self, model: torch.nn.Module):
         config = self.resolve_quantization_config()
         apply_quantization_config(model, config)
 
-        # apply observers, disable quantization until calibration
-        model.apply(self._initialize_observers)
+        # disable quantization until calibration
         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
+        Attach observers, 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(self._initialize_observers)
         model.apply(apply_calibration_status)
         model.apply(enable_quantization)  # quantize at the same time as calibrate
 
     def end_calibration(self, model: torch.nn.Module):
         """
-        Remove calibration hooks and set the model status to frozen. Keep quantization
-        enabled for future operations
+        Remove calibration hooks and observers, and set the model status to frozen.
+        Keep quantization enabled for future operations
 
         :param model: model to end calibration for
         """