# SPDX-FileCopyrightText: 2024-2025 ModelCloud.ai # SPDX-FileCopyrightText: 2024-2025 qubitium@modelcloud.ai # SPDX-License-Identifier: Apache-2.0 # Contact: qubitium@modelcloud.ai, x.com/qubitium # adapted from @qwopqwop200 's [GPTQ-for-LLaMa](https://github.com/qwopqwop200/GPTQ-for-LLaMa/tree/cuda), which itself is based on [gptq](https://github.com/IST-DASLab/gptq) import torch import torch.nn as nn from ..utils.logger import setup_logger from .config import BaseQuantizeConfig, _normalize_quant_bits, resolve_quant_format log = setup_logger() HF_OPTIMUM = "hf_optimum" def quantize(x, scale, zero, maxq, requires_groupwise_processing: bool): if maxq < 0: return (x > scale / 2).float() * scale + (x < zero / 2).float() * zero if requires_groupwise_processing: q = torch.clamp(torch.round(x / scale), -maxq, maxq) return scale * q else: q = torch.clamp(torch.round(x / scale) + zero, 0, maxq) return scale * (q - zero) class Quantizer(nn.Module): def __init__(self, qcfg: BaseQuantizeConfig, shape=1, name: str=None): super(Quantizer, self).__init__() self.qcfg = qcfg self.register_buffer("maxq", torch.tensor(0)) self.register_buffer("scale", torch.zeros(shape)) self.register_buffer("zero", torch.zeros(shape)) self.name=name def requires_groupwise_processing(self) -> bool: return False def _search_adjacent_zero_points(self) -> bool: """Return whether range search should optimize the 2-bit affine orientation.""" return ( int(self.maxq.item()) == 3 and not self.qcfg.sym and not self.requires_groupwise_processing() ) @staticmethod def _adjacent_zero_point( xmin: torch.Tensor, scale: torch.Tensor, primary: torch.Tensor, *, maxq: int, ) -> torch.Tensor: """Return the other integer neighboring the ideal affine zero point.""" ideal = -xmin / scale lower = torch.floor(ideal).clamp_(0, maxq) upper = torch.ceil(ideal).clamp_(0, maxq) return torch.where(primary == lower, upper, lower) # FIXME, optimum shouldn't call this directly, it should call hf_configure def configure( self, perchannel=False, grid=100, maxshrink=0.8, trits=False, bits: int | str | None = None, # for hf compat sym: bool | None = None, # for hf compat ): if self.name == HF_OPTIMUM: if bits is not None: self.qcfg.bits = _normalize_quant_bits(bits, format_value=resolve_quant_format(self.qcfg.format, self.qcfg.method)) if sym is not None: self.qcfg.sym = sym if self.requires_groupwise_processing(): self.maxq = torch.tensor(2 ** (self.qcfg.bits - 1) - 1) else: self.maxq = torch.tensor(2 ** self.qcfg.bits - 1) self.perchannel = perchannel self.grid = grid self.maxshrink = maxshrink if trits: self.maxq = torch.tensor(-1) def find_params(self, x, weight=False): dev = x.device self.maxq = self.maxq.to(dev) shape = x.shape if self.perchannel: if weight: x = x.flatten(1) else: if len(shape) == 4: x = x.permute([1, 0, 2, 3]) x = x.flatten(1) if len(shape) == 3: x = x.reshape((-1, shape[-1])).t() if len(shape) == 2: x = x.t() else: x = x.flatten().unsqueeze(0) tmp = torch.zeros(x.shape[0], device=dev) xmin = torch.minimum(x.min(1)[0], tmp) xmax = torch.maximum(x.max(1)[0], tmp) if self.qcfg.sym: xmax = torch.maximum(torch.abs(xmin), xmax) tmp = xmin < 0 if torch.any(tmp): xmin[tmp] = -xmax[tmp] tmp = (xmin == 0) & (xmax == 0) xmin[tmp] = -1 xmax[tmp] = +1 if self.maxq < 0: self.scale = xmax self.zero = xmin else: if self.requires_groupwise_processing(): self.scale = xmax / self.maxq self.zero = torch.zeros_like(self.scale) else: self.scale = (xmax - xmin) / self.maxq if self.qcfg.sym: self.zero = torch.full_like(self.scale, (self.maxq + 1) / 2) else: self.zero = torch.round(-xmin / self.scale) mse = float(getattr(self.qcfg, "mse", 0.0) or 0.0) if mse > 0.0: best = torch.full([x.shape[0]], float("inf"), device=dev) search_adjacent_zero_points = self._search_adjacent_zero_points() maxq_value = int(self.maxq.item()) if search_adjacent_zero_points else 0 for i in range(int(self.maxshrink * self.grid)): p = 1 - i / self.grid xmin1 = p * xmin xmax1 = p * xmax scale1 = ( xmax1 / self.maxq if self.requires_groupwise_processing() else (xmax1 - xmin1) / self.maxq ) zero1 = torch.round(-xmin1 / scale1) if not self.qcfg.sym else self.zero zero_candidates = (zero1,) if search_adjacent_zero_points: alternate_zero = self._adjacent_zero_point( xmin1, scale1, zero1, maxq=maxq_value, ) zero_candidates = (zero1, alternate_zero) for zero_candidate in zero_candidates: q = quantize( x, scale1.unsqueeze(1), zero_candidate.unsqueeze(1), self.maxq, self.requires_groupwise_processing(), ) q -= x q.abs_() q.pow_(mse) err = torch.sum(q, 1) tmp = err < best if torch.any(tmp): best[tmp] = err[tmp] self.scale[tmp] = scale1[tmp] self.zero[tmp] = zero_candidate[tmp] if not self.perchannel: if weight: tmp = shape[0] else: tmp = shape[1] if len(shape) != 3 else shape[2] self.scale = self.scale.repeat(tmp) self.zero = self.zero.repeat(tmp) if weight: shape = [-1] + [1] * (len(shape) - 1) self.scale = self.scale.reshape(shape) self.zero = self.zero.reshape(shape) return if len(shape) == 4: self.scale = self.scale.reshape((1, -1, 1, 1)) self.zero = self.zero.reshape((1, -1, 1, 1)) if len(shape) == 3: self.scale = self.scale.reshape((1, 1, -1)) self.zero = self.zero.reshape((1, 1, -1)) if len(shape) == 2: self.scale = self.scale.unsqueeze(0) self.zero = self.zero.unsqueeze(0) def quantize(self, x): return quantize(x, self.scale, self.zero, self.maxq, self.requires_groupwise_processing()) # def enabled(self): # return self.maxq > 0 # def ready(self): # return torch.all(self.scale != 0) class QQQQuantizer(Quantizer): def requires_groupwise_processing(self) -> bool: return self.qcfg.group_size == -1 and self.qcfg.sym __all__ = ["Quantizer"]