# 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 import math import threading import time from typing import Dict, Optional import torch import transformers from torch import nn from .. import QuantizeConfig from ..looper.named_module import NamedModule from ..quantization.config import FallbackStrategy, SmoothMSE from ..quantization.quantizer import HF_OPTIMUM from ..utils import setup_logger from ..utils.device import get_device from .fallback_smooth import mse_optimal_quant, smooth_block from .gptq import get_number_of_rows_and_cols from .npu_linalg import npu_inverse_cholesky_factor DEBUG = False log = setup_logger() def quantize(x, scale, zero, maxq, sym, groupsize): if maxq < 0: return (x > scale / 2).float() * scale + (x < zero / 2).float() * zero if groupsize != -1 or not sym: q = torch.clamp(torch.round(x / scale) + zero, 0, maxq) return scale * (q - zero) else: q = torch.clamp(torch.round(x / scale), -maxq, maxq) return scale * q class Quantizer(nn.Module): def __init__(self, shape=1): super(Quantizer, self).__init__() self.register_buffer("maxq", torch.tensor(0)) self.register_buffer("scale", torch.zeros(shape)) self.register_buffer("zero", torch.zeros(shape)) def configure( self, bits, perchannel=False, sym=True, groupsize=-1, mse=False, norm=2.4, grid=100, maxshrink=0.8, trits=False, ): if groupsize != -1 or not sym: self.maxq = torch.tensor(2 ** bits - 1) else: self.maxq = torch.tensor(2 ** (bits - 1) - 1) self.perchannel = perchannel self.groupsize = groupsize self.sym = sym self.mse = mse self.norm = norm 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.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.groupsize != -1 or not self.sym: self.scale = (xmax - xmin) / self.maxq if self.sym: self.zero = torch.full_like(self.scale, (self.maxq + 1) / 2) else: self.zero = torch.round(-xmin / self.scale) else: self.scale = xmax / self.maxq self.zero = torch.zeros_like(self.scale) if self.mse: best = torch.full([x.shape[0]], float("inf"), device=dev) for i in range(int(self.maxshrink * self.grid)): p = 1 - i / self.grid xmin1 = p * xmin xmax1 = p * xmax scale1 = ( (xmax1 - xmin1) / self.maxq if (self.groupsize != -1 or not self.sym) else xmax1 / self.maxq ) zero1 = torch.round(-xmin1 / scale1) if not self.sym else self.zero q = quantize( x, scale1.unsqueeze(1), zero1.unsqueeze(1), self.maxq, self.sym, self.groupsize, ) q -= x q.abs_() q.pow_(self.norm) err = torch.sum(q, 1) tmp = err < best if torch.any(tmp): best[tmp] = err[tmp] self.scale[tmp] = scale1[tmp] self.zero[tmp] = zero1[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): if self.ready(): return quantize( x, self.scale, self.zero, self.maxq, self.sym, self.groupsize ) return x def enabled(self): return self.maxq > 0 def ready(self): return torch.all(self.scale != 0) class QQQ: def __init__(self, module: nn.Module, qcfg: Optional[QuantizeConfig] = None): # self.lock = threading.Lock() # self.num_tied_handles = 0 # if qcfg.tied_gptq_handle is not None: # qcfg.tied_gptq_handle.num_tied_handles += 1 # Flags indicating issues # self.issue_zero_samples = False # self.issue_nan_hessian = False # self.issue_non_invertible = False # self.W = module.weight self.rows, self.columns = get_number_of_rows_and_cols(module) if isinstance(module, NamedModule): self.layer = module.module self.name = module.name self._named_module = module else: self.name = HF_OPTIMUM self.layer = module self._named_module = None layer_device = self.layer.weight.device setattr(self.layer, "target_device", layer_device) self.dev = self.layer.weight.device self._validate_module(self.layer) self.qcfg = qcfg if qcfg else QuantizeConfig() # HF compat will not pass qcfg self._original_rows = self.rows self._original_columns = self.columns if self._named_module is not None: pad_info = self._named_module.state.get("tp_pad_info") else: pad_info = getattr(self.layer, "_tp_pad_info", None) if isinstance(pad_info, dict): pad_cols = int(pad_info.get("pad_cols", 0) or 0) pad_cols = max(pad_cols, 0) else: pad_cols = 0 self._tp_pad_cols = pad_cols if self._tp_pad_cols: self.columns += self._tp_pad_cols self.module_copy = None self.nsamples = 0 self.quantizer = Quantizer() # fwd counter self.fwd_counter = 0 self.fallback = self.qcfg.fallback self.expected_nsamples: Optional[float] = None self.lock = threading.Lock() self.H: Optional[torch.Tensor] = None self._device_hessian_partials: Dict[torch.device, torch.Tensor] = {} self._device_sample_counts: Dict[torch.device, int] = {} @staticmethod def _validate_module(module): assert isinstance(module, (nn.Linear, nn.Conv1d, nn.Conv2d, transformers.Conv1D)), f"We supports only linear and convolutional layers. actual = `{module}`" @staticmethod def _truncate_last_dim(tensor: torch.Tensor, length: int) -> torch.Tensor: if tensor.dim() == 0: return tensor trim = min(length, tensor.shape[-1]) if trim == tensor.shape[-1]: return tensor return tensor.narrow(tensor.dim() - 1, 0, trim).contiguous() def _fallback_quantize(self, strategy: FallbackStrategy): maxq = 2 ** self.qcfg.bits - 1 sigma = 3.0 group_size = self.qcfg.group_size if self.qcfg.group_size != -1 else self.columns start_time = time.time() smooth_method = getattr(self.fallback, "smooth", None) mse_steps = 32 mse_maxshrink = 0.8 if isinstance(smooth_method, SmoothMSE): mse_steps = smooth_method.steps mse_maxshrink = smooth_method.maxshrink W = self.layer.weight.data.clone() if isinstance(self.layer, nn.Conv2d): W = W.flatten(1) if isinstance(self.layer, transformers.Conv1D): W = W.t() W = W.float() if self._tp_pad_cols: pad = torch.zeros( (W.shape[0], self._tp_pad_cols), dtype=W.dtype, device=W.device, ) W = torch.cat((W, pad), dim=1) Q = torch.empty_like(W) scale_chunks = [] zero_chunks = [] for start in range(0, self.columns, group_size): end = min(start + group_size, self.columns) block = W[:, start:end] if isinstance(smooth_method, SmoothMSE): dequant, scale, zero = mse_optimal_quant( block, self.qcfg, maxq, steps=mse_steps, maxshrink=mse_maxshrink, ) else: block_mod, scale_factor = smooth_block( block, self.fallback, group_size=self.qcfg.group_size if self.qcfg.group_size != -1 else self.columns, ) if strategy == FallbackStrategy.MIDPOINT: w_min = block_mod.min(dim=1, keepdim=True).values w_max = block_mod.max(dim=1, keepdim=True).values mid = (w_max + w_min) / 2.0 scale = torch.clamp((w_max - w_min) / maxq, min=1e-8) zero_mid = torch.full_like(scale, maxq / 2.0) q = torch.round((block_mod - mid) / scale + zero_mid) q = torch.clamp(q, 0, maxq) zero = torch.round(zero_mid - (mid / scale)) zero = torch.clamp(zero, 0, maxq) dequant = (q - zero) * scale elif strategy == FallbackStrategy.MEAN: mean = block_mod.mean(dim=1, keepdim=True) max_dev = torch.max((block_mod - mean).abs(), dim=1, keepdim=True).values max_dev = torch.clamp(max_dev, min=1e-8) scale = (2 * max_dev) / maxq zero_mid = torch.full_like(scale, maxq / 2.0) q = torch.round((block_mod - mean) / scale + zero_mid) q = torch.clamp(q, 0, maxq) zero = torch.round(zero_mid - (mean / scale)) zero = torch.clamp(zero, 0, maxq) dequant = (q - zero) * scale elif strategy == FallbackStrategy.MEDIAN: median = block_mod.median(dim=1, keepdim=True).values max_dev = torch.max((block_mod - median).abs(), dim=1, keepdim=True).values max_dev = torch.clamp(max_dev, min=1e-8) scale = (2 * max_dev) / maxq zero_mid = torch.full_like(scale, maxq / 2.0) q = torch.round((block_mod - median) / scale + zero_mid) q = torch.clamp(q, 0, maxq) zero = torch.round(zero_mid - (median / scale)) zero = torch.clamp(zero, 0, maxq) dequant = (q - zero) * scale elif strategy == FallbackStrategy.STDCLIP: mean = block_mod.mean(dim=1, keepdim=True) std = block_mod.std(dim=1, keepdim=True, unbiased=False) std = torch.clamp(std, min=1e-8) lo = mean - sigma * std hi = mean + sigma * std scale = torch.clamp((hi - lo) / maxq, min=1e-8) zero = torch.round(-lo / scale) zero = torch.clamp(zero, 0, maxq) q = torch.round(block_mod / scale + zero) q = torch.clamp(q, 0, maxq) dequant = (q - zero) * scale elif strategy == FallbackStrategy.RTN: self.quantizer.find_params(block_mod, weight=True) dequant = self.quantizer.quantize(block_mod) scale = self.quantizer.scale zero = self.quantizer.zero else: raise ValueError(f"Unsupported fallback strategy: {strategy}") if scale_factor is not None: scale = scale * scale_factor dequant = dequant * scale_factor Q[:, start:end] = dequant scale_block = scale if scale.dim() > 1 else scale.unsqueeze(1) zero_block = zero if zero.dim() > 1 else zero.unsqueeze(1) if scale_block.shape[1] > 1: scale_block = scale_block.mean(dim=1, keepdim=True) if zero_block.shape[1] > 1: zero_block = zero_block.mean(dim=1, keepdim=True) scale_chunks.append(scale_block) zero_chunks.append(zero_block) scale = torch.cat(scale_chunks, dim=1) zero = torch.cat(zero_chunks, dim=1) if self._tp_pad_cols: valid_cols = self._original_columns Q = Q[:, :valid_cols] scale = self._truncate_last_dim(scale, valid_cols) zero = self._truncate_last_dim(zero, valid_cols) else: valid_cols = self.columns g_idx = torch.arange(valid_cols, device=Q.device, dtype=torch.int32) // (group_size if group_size != -1 else self.columns) if isinstance(self.layer, transformers.Conv1D): Q = Q.t() if Q.shape != self.layer.weight.shape: Q = Q.reshape(self.layer.weight.shape).to(self.layer.weight.dtype) else: Q = Q.to(self.layer.weight.dtype) scale_extra = None if group_size != self.columns: quantizer_extra = Quantizer() quantizer_extra.configure( bits=8, perchannel=True, groupsize=-1, sym=True, mse=False, ) quantizer_extra.find_params(self.layer.weight.data.clone(), weight=True) scale_extra = quantizer_extra.scale duration = time.time() - start_time mean_abs_err = (Q - self.layer.weight.data).abs().mean().item() avg_loss = f"fallback({strategy.value}): {mean_abs_err:.7f}" damp_percent = 0.0 self.H = None return Q, scale, zero, g_idx, duration, avg_loss, damp_percent, scale_extra, self.nsamples def add_batch(self, inp, out): if DEBUG: self.inp1 = inp self.out1 = out if len(inp.shape) == 2: inp = inp.unsqueeze(0) if isinstance(self.layer, nn.Linear) or isinstance( self.layer, transformers.Conv1D ): if len(inp.shape) == 3: inp = inp.reshape((-1, inp.shape[-1])) inp = inp.t() if isinstance(self.layer, nn.Conv2d): unfold = nn.Unfold( self.layer.kernel_size, dilation=self.layer.dilation, padding=self.layer.padding, stride=self.layer.stride, ) inp = unfold(inp) inp = inp.permute([1, 0, 2]) inp = inp.flatten(1) dev = torch.device(get_device(inp)) batch_token_size = inp.shape[1] inp = inp.float() if self._tp_pad_cols: pad = inp.new_zeros((self._tp_pad_cols, inp.shape[1])) inp = torch.cat((inp, pad), dim=0) with self.lock: self.fwd_counter += 1 existing = self._device_hessian_partials.get(dev) if existing is None: existing = torch.zeros( (self.columns, self.columns), dtype=torch.float32, device=dev, ) self._device_hessian_partials[dev] = existing # Accumulate xtx directly into the per-device buffer to avoid # allocating a separate output tensor for every batch. existing.addmm_(inp, inp.t(), beta=1.0, alpha=1.0) self._device_sample_counts[dev] = self._device_sample_counts.get(dev, 0) + batch_token_size self.nsamples += batch_token_size def materialize_hessian(self, target_device: Optional[torch.device] = None) -> torch.Tensor: if target_device is None: target_device = self.dev target_device = torch.device(target_device) total_samples = sum(self._device_sample_counts.values()) if total_samples == 0: self.H = torch.zeros((self.columns, self.columns), dtype=torch.float32, device=target_device) self.nsamples = 0 return self.H # Merge per-device partials into a single Hessian on the target device. # Pop and delete each partial as it is added so the peak memory stays # at result + one partial instead of result + all partials. result = torch.zeros((self.columns, self.columns), dtype=torch.float32, device=target_device) while self._device_hessian_partials: dev, partial = self._device_hessian_partials.popitem() if partial.device == target_device and partial.dtype == torch.float32: result.add_(partial) else: result.add_(partial.to(device=target_device, dtype=torch.float32)) del partial result.mul_(2.0 / float(total_samples)) self.H = result self.nsamples = total_samples self._device_sample_counts.clear() return self.H @torch.inference_mode() def quantize( self, blocksize=128, ): start = time.time() from ..utils.fallback import resolve_fallback_strategy, resolve_threshold, should_use_fallback resolved_strategy = resolve_fallback_strategy(self.fallback) self.materialize_hessian() percdamp = self.qcfg.damp_percent groupsize = self.qcfg.group_size actorder = self.qcfg.desc_act static_groups = self.qcfg.static_groups W = self.layer.weight.data.clone() if isinstance(self.layer, nn.Conv2d): W = W.flatten(1) if isinstance(self.layer, transformers.Conv1D): W = W.t() W = W.float() if self._tp_pad_cols: pad = torch.zeros( (W.shape[0], self._tp_pad_cols), dtype=W.dtype, device=W.device, ) W = torch.cat((W, pad), dim=1) tick = time.time() if not self.quantizer.ready(): self.quantizer.find_params(W, weight=True) H = self.H del self.H dead = torch.diag(H) == 0 H[dead, dead] = 1 W[:, dead] = 0 g_idx = [] scale = [] zero = [] now_idx = 1 if static_groups: import copy groups = [] for i in range(0, self.columns, groupsize): quantizer = copy.deepcopy(self.quantizer) quantizer.find_params(W[:, i: (i + groupsize)], weight=True) scale.append(quantizer.scale) zero.append(quantizer.zero) groups.append(quantizer) if actorder: perm = torch.argsort(torch.diag(H), descending=True) W = W[:, perm] H = H[perm][:, perm] invperm = torch.argsort(perm) Losses = torch.zeros_like(W) Q = torch.zeros_like(W) damp = percdamp * torch.mean(torch.diag(H)) diag = torch.arange(self.columns, device=H.device) H[diag, diag] += damp threshold_raw, is_percent = resolve_threshold(self.fallback, self.expected_nsamples) fallback_configured = threshold_raw is not None try: if H.device.type == "npu": H = npu_inverse_cholesky_factor(H) else: H = torch.linalg.cholesky(H) H = torch.cholesky_inverse(H) H = torch.linalg.cholesky(H, upper=True) Hinv = H except Exception: fallback_requested = should_use_fallback( self.fallback, float(self.nsamples), self.expected_nsamples, ) if fallback_requested: extra = f", threshold_raw={threshold_raw}" if threshold_raw is not None and is_percent else "" log.warn( "Quantization: Module `%s` -> Using `%s` fallback quantization (observed %s samples, threshold=%s%s, max_total=%s).", self.name, resolved_strategy.value, self.nsamples, self.fallback, extra, self.expected_nsamples, ) if resolved_strategy != FallbackStrategy.RTN: return self._fallback_quantize(resolved_strategy) Hinv = None else: raise for i1 in range(0, self.columns, blocksize): i2 = min(i1 + blocksize, self.columns) count = i2 - i1 W1 = W[:, i1:i2].clone() Q1 = torch.zeros_like(W1) Err1 = torch.zeros_like(W1) Losses1 = torch.zeros_like(W1) Hinv1 = Hinv[i1:i2, i1:i2] for i in range(count): w = W1[:, i] d = Hinv1[i, i] if groupsize != -1: if not static_groups: if (i1 + i) % groupsize == 0: self.quantizer.find_params( W[:, (i1 + i): (i1 + i + groupsize)], weight=True ) if ((i1 + i) // groupsize) - now_idx == -1: scale.append(self.quantizer.scale) zero.append(self.quantizer.zero) now_idx += 1 else: idx = i1 + i if actorder: idx = perm[idx] self.quantizer = groups[idx // groupsize] q = quantize( w.unsqueeze(1), self.quantizer.scale, self.quantizer.zero, self.quantizer.maxq, self.quantizer.sym, self.quantizer.groupsize, ).flatten() Q1[:, i] = q Losses1[:, i] = (w - q) ** 2 / d ** 2 err1 = (w - q) / d W1[:, i:] -= err1.unsqueeze(1).matmul(Hinv1[i, i:].unsqueeze(0)) Err1[:, i] = err1 Q[:, i1:i2] = Q1 Losses[:, i1:i2] = Losses1 / 2 W[:, i2:] -= Err1.matmul(Hinv[i1:i2, i2:]) if DEBUG: self.layer.weight.data[:, :i2] = Q[:, :i2] self.layer.weight.data[:, i2:] = W[:, i2:] print(torch.sum((self.layer(self.inp1) - self.out1) ** 2)) print(torch.sum(Losses)) torch.cuda.synchronize() print("time %.2f" % (time.time() - tick)) print("error", torch.sum(Losses).item()) if Hinv is not None: del Hinv if self.nsamples != 0: avg_loss = torch.sum(Losses).item() / self.nsamples if math.isnan(avg_loss): print("Losses sum item:", torch.sum(Losses).item()) if fallback_configured: log.info(f"Quantization: Failed due to `NaN` loss for `{self.name}`, use mock quantization retry for `{self.name}`") self.qcfg.mock_quantization = True return self.quantize(blocksize=blocksize) else: raise ValueError(f"Quantization: Failed due to `NaN` loss for `{self.name}`, please try increasing calibration data samples or enable fallback=True") else: if fallback_configured: log.warn(f"Quantization: Module `{self.name}` -> using fail safe mode. Please check if calibration data is sufficient.") else: log.warn(f"Quantization: `{self.name}` is not activated due to model inference logic (MoE)") avg_loss = f"{resolved_strategy.value} fallback" if fallback_configured else 999999999 else: avg_loss = f"{resolved_strategy.value} fallback" if fallback_configured else 999999999 del Losses groupsize = groupsize if groupsize != -1 else self.columns if static_groups and actorder: g_idx = [perm[i] // groupsize for i in range(self.columns)] else: g_idx = [i // groupsize for i in range(self.columns)] g_idx = torch.tensor(g_idx, dtype=torch.int32, device=Q.device) if actorder: Q = Q[:, invperm] g_idx = g_idx[invperm] if self._tp_pad_cols: valid_cols = self._original_columns Q = Q[:, :valid_cols] g_idx = g_idx[:valid_cols] if isinstance(self.layer, transformers.Conv1D): Q = Q.t() if Q.shape != self.layer.weight.shape: Q = Q.reshape(self.layer.weight.shape).to(self.layer.weight.dtype) else: Q = Q.to(self.layer.weight.dtype) if DEBUG: print(torch.sum((self.layer(self.inp1) - self.out1) ** 2)) if scale == []: scale.append(self.quantizer.scale) zero.append(self.quantizer.zero) scale = torch.cat(scale, dim=1) zero = torch.cat(zero, dim=1) if self._tp_pad_cols: valid_cols = self._original_columns scale = self._truncate_last_dim(scale, valid_cols) zero = self._truncate_last_dim(zero, valid_cols) Q = Q.to(device=self.layer.weight.data.device, non_blocking=False) # post int8 quant scale_extra = None if groupsize != self.columns: quantizer_extra = Quantizer() quantizer_extra.configure( bits=8, perchannel=True, groupsize=-1, sym=True, mse=False, ) quantizer_extra.find_params(self.layer.weight.data.clone(), weight=True) scale_extra = quantizer_extra.scale duration = time.time() - start return Q, scale, zero, g_idx, duration, avg_loss, damp, scale_extra, self.nsamples def free(self): if DEBUG: self.inp1 = None self.out1 = None self.H = None self._device_hessian_partials.clear() self._device_sample_counts.clear() self.Losses = None self.Trace = None if hasattr(self, "quantizer"): del self.quantizer if hasattr(self, "module_copy"): del self.module_copy if self._named_module is not None: self._named_module.state.pop("tp_pad_info", None) torch.cuda.empty_cache()