llmcompressor.observers

@classmethod
def fuse(cls, observers_and_modules: Iterable[tuple["Observer", Module]]) -> None:
    """
    Link all observers in the iterable with each other for shared global_scale.
    Sets module weak-references and makes handlers aware of each other.

    :param observers_and_modules: iterable of (observer, module) tuples
    """
    pairs = list(observers_and_modules)
    handlers = []
    for obs, mod in pairs:
        obs.fusion_handler._module_ref = ref(mod)
        handlers.append(obs.fusion_handler)

    for handler in handlers:
        handler._group = [h for h in handlers if h is not handler]

def get_fused_statistics(self) -> list[dict[str, torch.Tensor]]:
    """
    Returns statistics for every observer in the fusion group
    (including this handler's own observer).

    If a fused observer lacks statistics, runs observation on its
    module's weight via the weak ref stored during fuse().

    :return: list of {"min_vals": Tensor, "max_vals": Tensor} dicts
    """
    all_handlers = [self] + self._group

    for handler in all_handlers:
        obs = handler._observer
        if not obs.has_statistics:
            mod = handler.module
            assert mod is not None, _msg
            obs(mod.weight)

    return [h._observer.get_statistics() for h in all_handlers]

def maybe_delete_statistics(self) -> None:
    """
    Cooperative statistics deletion. Each handler calls this when it
    is done with fused statistics. When ALL handlers in the group
    have called it, deletes statistics from every observer in the
    group and resets the bookkeeping.

    For unfused observers, deletes statistics immediately.
    """
    if not self.is_fused:
        self._delete_observer_statistics(self._observer)
        return

    self._deletion_called = True

    if all(h._deletion_called for h in self._group):
        self._deletion_called = False
        for h in self._group:
            h._deletion_called = False
            self._delete_observer_statistics(h._observer)
        self._delete_observer_statistics(self._observer)

def attach(self, module: torch.nn.Module) -> None:
    """Attach a forward-pre hook to accumulate E[x²] per input channel.

    If raw accumulators (``_imatrix_sum`` / ``_imatrix_count``) already
    exist on the module (second pass after IMatrixGatherer), copy them
    to the observer and skip hook registration.
    """
    if hasattr(module, "_imatrix_sum"):
        self._imatrix_sum = module._imatrix_sum
        self._imatrix_count = module._imatrix_count
        del module._imatrix_sum
        del module._imatrix_count
        return

    if not hasattr(module, "in_features"):
        return

    in_features = module.in_features
    module._imatrix_sum = torch.zeros(in_features, dtype=IMATRIX_PRECISION)
    module._imatrix_count = torch.tensor(0, dtype=torch.int64)

    def _hook(mod, args):
        if (
            HooksMixin._HOOKS_DISABLED
            and getattr(mod, "_imatrix_hook", None)
            not in HooksMixin._HOOKS_KEEP_ENABLED
        ):
            return
        x = args[0] if isinstance(args, tuple) else args
        if isinstance(x, tuple):
            x = x[0]
        if x is None or not isinstance(x, torch.Tensor):
            return

        x_f = x.detach().to(IMATRIX_PRECISION)
        device = x_f.device
        n_tokens = math.prod(x_f.shape[:-1])
        token_sum = x_f.pow(2).sum(dim=list(range(x_f.dim() - 1)))

        mod._imatrix_sum = mod._imatrix_sum.to(device)
        mod._imatrix_count = mod._imatrix_count.to(device)

        mod._imatrix_sum.add_(token_sum)
        mod._imatrix_count += n_tokens

    module._imatrix_hook = module.register_forward_pre_hook(_hook)

def detach(self, module: torch.nn.Module) -> None:
    """Remove hooks and leave raw sum/count on module for second-pass pickup.

    Case 1 – accumulators present on module: leave them for next
    observer's ``attach()`` to pick up.

    Case 2 – no accumulators (second-pass cleanup): nothing to do.
    """
    if hasattr(module, "_imatrix_hook"):
        module._imatrix_hook.remove()
        del module._imatrix_hook

def attach(self, module: torch.nn.Module) -> None:
    """
    Called when the observer is attached to a module.
    Subclasses can override to register hooks or initialize state.

    :param module: the module this observer is being attached to
    """
    pass

def delete_statistics(self, check_fused: bool = True) -> None:
    """
    Delete this observer's statistics.

    :param check_fused: if True (default), use cooperative fusion deletion
        so stats are only removed when all fused observers are ready.
        If False, delete immediately regardless of fusion state.
    """
    if check_fused:
        self.fusion_handler.maybe_delete_statistics()
    else:
        for attr in self._stats_attrs:
            if hasattr(self, attr):
                delattr(self, attr)

def detach(self, module: torch.nn.Module) -> None:
    """
    Called before the observer is deleted from a module.
    Subclasses can override to remove hooks and clean up module attributes.

    :param module: the module this observer is being removed from
    """
    pass

@torch.no_grad
def forward(self, observed: torch.Tensor) -> "Observer":
    """
    Update observer statistics from observed value.
    If fused, also triggers observation for all fused observers.

    :param observed: value being observed
    :return: self for method chaining
    """
    if observed.numel() == 0:
        return self

    if self.is_weight_obs and self.has_statistics:
        return self

    observed = flatten_for_calibration(observed, self.base_name, self.args)
    self.update_statistics_from_observed(observed)

    self.fusion_handler.get_fused_statistics()

    return self

def get_global_scale(self) -> Optional[torch.Tensor]:
    """
    Compute global_scale from this observer and all fused observers.

    Uses the FusionHandler to collect statistics from the fusion group.
    Returns None if strategy is not TENSOR_GROUP.

    :return: global_scale tensor or None
    """
    if self.args.strategy != QuantizationStrategy.TENSOR_GROUP:
        return None

    all_stats = self.fusion_handler.get_fused_statistics()
    global_absmax = all_stats[0]["max_vals"].max()
    for stats in all_stats:
        global_absmax = torch.max(global_absmax, -stats["min_vals"].min())
        global_absmax = torch.max(global_absmax, stats["max_vals"].max())

    return generate_gparam(-global_absmax.reshape(1), global_absmax.reshape(1))

@torch.no_grad
def get_qparams(self) -> QParamsDict:
    """
    Compute quantization parameters from accumulated statistics.

    For TENSOR_GROUP, global_scale is computed from the absmax of
    this observer and all fused observers.

    :return: dict with keys "scale", "zero_point", and "global_scale"
    """
    assert (
        self.has_statistics
    ), "No statistics available. Call observer(value) first."

    global_scale = self.get_global_scale()

    scale, zero_point = calculate_qparams(
        min_vals=self.min_vals,
        max_vals=self.max_vals,
        quantization_args=self.args,
        global_scale=global_scale,
    )

    self.delete_statistics()

    return {"scale": scale, "zero_point": zero_point, "global_scale": global_scale}

def get_statistics(self) -> dict[str, torch.Tensor]:
    """
    Return this observer's current statistics.

    :return: dict mapping each stats attr name to its tensor
    """
    assert (
        self.has_statistics
    ), "No statistics available. Call observer(value) first."
    return {attr: getattr(self, attr) for attr in self._stats_attrs}

def sync_activation_stats(self) -> List[dist.Work]:
    """All-reduce accumulated activation statistics across DDP ranks.

        note: weight statistics don't need to be synced since weights
    are synced across ranks, only data (activations) differs by rank.

    :return: list of async communication handles
    """
    comms = []
    for attr_name, reduce_op in self._act_sync_dict.items():
        val = getattr(self, attr_name, None)
        if val is not None:
            comms.append(
                dist.all_reduce(as_broadcastable(val), op=reduce_op, async_op=True)
            )
    return comms

@abstractmethod
def update_statistics_from_observed(self, observed: torch.Tensor) -> None:
    """
    Update internal observer statistics (min_vals, max_vals) from observed tensor.

    :param observed: flattened observed value of shape
                    (num_observations, *qparam_shape, group_size)
    """
    raise NotImplementedError()