feat: defer activation qparam calculation to sequential epoch end

Fixes #2446

Author: dzhengAP

src/llmcompressor/modifiers/quantization/calibration.py

@@ -17,6 +17,7 @@
 from torch.nn import Module
 
 from llmcompressor.observers import Observer
+from llmcompressor.observers.base import calibrate_module_from_observer
 
 __all__ = [
     "initialize_observer",
@@ -30,6 +31,7 @@
     "calibrate_query_hook",
     "calibrate_key_hook",
     "calibrate_value_hook",
+    "flush_activation_qparams",
 ]
 
 
@@ -156,15 +158,20 @@ def update_weight_zp_scale(module: Module):
     call_observer(module=module, base_name="weight")
 
 
-def calibrate_activations(module: Module, value: torch.Tensor, base_name: str):
+def calibrate_activations(
+    module: Module, value: torch.Tensor, base_name: str, stats_only: bool = False
+):
     """
     Calibrate input or output activations by calling the a module's attached
     observer.
 
     :param module: torch.nn.Module
     :param base_name: substring used to fetch the observer, scales, and zp
     :param value: torch.Tensor to be passed to the observer
-
+    :param stats_only: if True, only update running statistics in the observer
+        (accumulate min/max) without computing or writing scale/zero_point.
+        Used during deferred qparam calibration — qparams are computed once
+        at epoch end via flush_activation_qparams instead of per batch.
     """
     # If empty tensor, can't update zp/scale
     # Case for MoEs
@@ -184,6 +191,12 @@ def calibrate_activations(module: Module, value: torch.Tensor, base_name: str):
         if quantization_args.strategy == QuantizationStrategy.TENSOR_GROUP:
             calculate_gparam = True
 
+    # In deferred (stats_only) mode, only accumulate running min/max in the
+    # observer — skip writing scale/zero_point until epoch end.
+    if stats_only:
+        calculate_qparams = False
+        calculate_gparam = False
+
     call_observer(
         module=module,
         base_name=base_name,
@@ -196,43 +209,40 @@ def calibrate_activations(module: Module, value: torch.Tensor, base_name: str):
 def calibrate_input_hook(module: Module, args: Any):
     """
     Hook to calibrate input activations.
-    Will call the observers to update the scales/zp before applying
-    input QDQ in the module's forward pass.
+    Accumulates running min/max statistics in the observer without computing
+    scale/zero_point. Qparams are computed once at epoch end via
+    flush_activation_qparams (deferred mode).
     """
     args = args[0] if isinstance(args, tuple) else args
-    calibrate_activations(module, value=args, base_name="input")
+    calibrate_activations(module, value=args, base_name="input", stats_only=True)
 
 
 def calibrate_output_hook(module: Module, _args: Any, output: torch.Tensor):
     """
     Hook to calibrate output activations.
-    Will call the observers to update the scales/zp before applying
-    output QDQ.
+    Accumulates running min/max statistics only (deferred qparam mode).
+    Qparams are computed at epoch end; forward_quantize is skipped during
+    calibration batches since quantization is disabled in the sequential pipeline.
     """
     calibrate_activations(
         module,
         value=output,
         base_name="output",
-    )
-    output = forward_quantize(
-        module=module,
-        value=output,
-        base_name="output",
-        args=module.quantization_scheme.output_activations,
+        stats_only=True,
     )
     return output
 
 
 def calibrate_query_hook(module: Module, query_states: torch.Tensor):
-    calibrate_activations(module, query_states, base_name="q")
+    calibrate_activations(module, query_states, base_name="q", stats_only=True)
 
 
 def calibrate_key_hook(module: Module, key_states: torch.Tensor):
-    calibrate_activations(module, key_states, base_name="k")
+    calibrate_activations(module, key_states, base_name="k", stats_only=True)
 
 
 def calibrate_value_hook(module: Module, value_states: torch.Tensor):
-    calibrate_activations(module, value_states, base_name="v")
+    calibrate_activations(module, value_states, base_name="v", stats_only=True)
 
 
 def apply_calibration_status(module: Module):
@@ -273,3 +283,29 @@ def reset_quantization_status(model: Module):
     for module in model.modules():
         if hasattr(module, "quantization_status"):
             delattr(module, "quantization_status")
+
+
+def flush_activation_qparams(module: Module):
+    """
+    Compute and write final activation qparams from each observer's accumulated
+    running statistics, then free those statistics to reduce memory.
+
+    This is called once at SEQUENTIAL_EPOCH_END for each subgraph, replacing the
+    per-batch qparam updates that were previously triggered by calibration hooks.
+    It is a no-op for modules with no quantization scheme or no activation observers.
+
+    Note: weight observers are not touched here — weight qparams are always computed
+    up-front in ``on_start`` via ``update_weight_zp_scale``.
+
+    apply to targeted modules with:
+        for _, module in match_named_modules(...):
+            flush_activation_qparams(module)
+
+    :param module: module to flush activation qparams for
+    """
+    scheme = getattr(module, "quantization_scheme", None)
+    if scheme is None:
+        return
+
+    for base_name in ("input", "output", "q", "k", "v"):
+        calibrate_module_from_observer(module, base_name)

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

@@ -4,6 +4,7 @@
 from llmcompressor.core import Event, EventType, State
 from llmcompressor.modifiers import Modifier
 from llmcompressor.modifiers.quantization.calibration import (
+    flush_activation_qparams,
     update_weight_global_scale,
     update_weight_zp_scale,
 )
@@ -65,7 +66,10 @@ def on_initialize(self, state: State, **kwargs) -> bool:
 
     def on_start(self, state: State, event: Event, **kwargs):
         """
-        Begin calibrating activations and weights. Calibrate weights only once on start
+        Begin calibrating activations and weights. Calibrate weights only once on start.
+        Quantization is kept DISABLED during calibration batches so that forward passes
+        run in fp32. Activation qparams are computed once per subgraph at
+        SEQUENTIAL_EPOCH_END via flush_activation_qparams (deferred mode).
         """
         self.started_ = True
         QuantizationMixin.start_calibration(self, state.model)
@@ -90,11 +94,26 @@ def on_start(self, state: State, event: Event, **kwargs):
         for _, module in tqdm.tqdm(named_modules, desc="Calibrating weights"):
             update_weight_zp_scale(module)
 
+        # Disable quantization during calibration batches so that fp32 activations
+        # flow through the model unmodified while hooks accumulate running stats.
+        # Re-enable once after epoch end when qparams have been flushed.
+        from compressed_tensors.quantization import disable_quantization
+
+        state.model.apply(disable_quantization)
+
     def on_event(self, state: State, event: Event, **kwargs):
         if event.type_ == EventType.CALIBRATION_EPOCH_START:
             if not self.started_:
                 self.on_start(state, None)
 
+        if event.type_ == EventType.SEQUENTIAL_EPOCH_END:
+            # Deferred qparam flush: compute scale/zero_point from accumulated
+            # running statistics, then free those stats to reduce memory.
+            for _, module in match_named_modules(
+                state.model, self.resolved_targets, self.ignore
+            ):
+                flush_activation_qparams(module)
+
         if event.type_ == EventType.CALIBRATION_EPOCH_END:
             if not self.ended_:
                 self.on_end(state, None)

src/llmcompressor/observers/base.py

@@ -11,7 +11,7 @@
 
 from llmcompressor.observers.helpers import flatten_for_calibration
 
-__all__ = ["Observer", "MinMaxTuple", "ScaleZpTuple"]
+__all__ = ["Observer", "MinMaxTuple", "ScaleZpTuple", "calibrate_module_from_observer"]
 
 MinMaxTuple = Tuple[torch.Tensor, torch.Tensor]
 ScaleZpTuple = Tuple[torch.Tensor, torch.Tensor]
@@ -74,6 +74,39 @@ def get_global_min_max(self, observed: torch.Tensor) -> MinMaxTuple:
         """
         raise NotImplementedError()
 
+    def get_accumulated_min_max(self) -> Optional[MinMaxTuple]:
+        """
+        Return the accumulated running min/max statistics stored by this observer,
+        without performing any new observation. Returns None if no statistics have
+        been accumulated yet (i.e. no batches have been seen).
+
+        Subclasses which accumulate state (StaticMinMax, MovingAverage) naturally
+        expose this through their ``past_min_vals`` / ``past_max_vals`` attributes.
+        Memoryless observers have no running state, so this always returns None.
+
+        :return: (min_vals, max_vals) tensors or None
+        """
+        min_vals = getattr(self, "past_min_vals", None)
+        max_vals = getattr(self, "past_max_vals", None)
+        if min_vals is None or max_vals is None:
+            return None
+        return min_vals, max_vals
+
+    def clear_accumulated_stats(self):
+        """
+        Delete accumulated running statistics to free memory after qparams have been
+        computed and written to the parent module. Only clears attributes that exist
+        on the observer (memoryless observers are unaffected).
+        """
+        for attr in (
+            "past_min_vals",
+            "past_max_vals",
+            "past_global_min_vals",
+            "past_global_max_vals",
+        ):
+            if hasattr(self, attr):
+                delattr(self, attr)
+
     @torch.no_grad
     def forward(self, observed: torch.Tensor) -> ScaleZpTuple:
         """
@@ -142,3 +175,50 @@ def _check_has_global_scale(self, global_scale: Optional[torch.nn.Parameter]):
                 "Cannot compute scale and zero points "
                 "without first computing global scale"
             )
+
+
+@torch.no_grad()
+def calibrate_module_from_observer(
+    module: torch.nn.Module,
+    base_name: str,
+) -> bool:
+    """
+    Flush an observer's accumulated running statistics into the parent module's
+    quantization parameters (scale / zero_point), then free the running stats.
+
+    This is the deferred counterpart to ``call_observer``. Instead of accepting a
+    fresh activation tensor, it reads the min/max values that the observer has
+    already accumulated across all calibration batches and computes qparams from
+    those final statistics.
+
+    :param module: module whose ``{base_name}_observer`` attribute holds the observer
+    :param base_name: one of "input", "output", "q", "k", "v"
+    :return: True if qparams were updated, False if observer had no accumulated stats
+    """
+    from compressed_tensors.utils import align_module_device, update_offload_parameter
+
+    observer: Optional[Observer] = getattr(module, f"{base_name}_observer", None)
+    if observer is None:
+        return False
+
+    accumulated = observer.get_accumulated_min_max()
+    if accumulated is None:
+        return False
+
+    min_vals, max_vals = accumulated
+    global_scale = getattr(module, f"{base_name}_global_scale", None)
+
+    with align_module_device(module):
+        scales, zero_points = calculate_qparams(
+            min_vals=min_vals,
+            max_vals=max_vals,
+            quantization_args=observer.args,
+            global_scale=global_scale,
+        )
+        update_offload_parameter(module, f"{base_name}_scale", scales)
+        if hasattr(module, f"{base_name}_zero_point"):
+            update_offload_parameter(module, f"{base_name}_zero_point", zero_points)
+
+    # Free memory — running stats no longer needed
+    observer.clear_accumulated_stats()
+    return True