Spaces:
Running on Zero
Running on Zero
| # quantization/cfrn/quantlib/quant_utils.py | |
| from __future__ import annotations | |
| import math | |
| from dataclasses import dataclass | |
| from typing import Optional, Tuple | |
| import torch | |
| # --------------------------- | |
| # Helper: int4 pack / unpack | |
| # --------------------------- | |
| def pack_int4(q_int8: torch.Tensor) -> Tuple[torch.Tensor, torch.Size]: | |
| """ | |
| Pack int4 values stored in int8 tensor (expected range [-8,7]) into uint8. | |
| Two int4 values per byte: low nibble is first, high nibble is second. | |
| Returns: packed_uint8, original_shape | |
| """ | |
| if q_int8.dtype != torch.int8: | |
| raise TypeError(f"pack_int4 expects int8, got {q_int8.dtype}") | |
| shape = q_int8.shape | |
| flat = q_int8.reshape(-1).to(torch.int16) # avoid overflow | |
| # map [-8,7] -> [0,15] | |
| flat = flat + 8 | |
| flat = torch.clamp(flat, 0, 15).to(torch.uint8) | |
| if flat.numel() % 2 == 1: | |
| flat = torch.cat([flat, torch.zeros(1, device=flat.device, dtype=flat.dtype)], dim=0) | |
| lo = flat[0::2] | |
| hi = flat[1::2] | |
| packed = lo | (hi << 4) | |
| return packed, shape | |
| def unpack_int4(packed_uint8: torch.Tensor, shape: torch.Size) -> torch.Tensor: | |
| """ | |
| Unpack uint8 -> int8 int4 values in [-8,7], restoring given shape. | |
| """ | |
| if packed_uint8.dtype != torch.uint8: | |
| raise TypeError(f"unpack_int4 expects uint8, got {packed_uint8.dtype}") | |
| flat = packed_uint8.reshape(-1) | |
| lo = flat & 0x0F | |
| hi = (flat >> 4) & 0x0F | |
| out = torch.empty((flat.numel() * 2,), device=flat.device, dtype=torch.int16) | |
| out[0::2] = lo.to(torch.int16) | |
| out[1::2] = hi.to(torch.int16) | |
| out = out[: int(torch.tensor(shape).prod().item())] | |
| out = out.to(torch.int16) - 8 | |
| return out.to(torch.int8).reshape(shape) | |
| # --------------------------- | |
| # Symmetric quant/dequant | |
| # --------------------------- | |
| class QParams: | |
| scale: torch.Tensor # float tensor, can be scalar or per-channel/group | |
| qmax: int | |
| qmin: int | |
| bits: int | |
| per_channel: bool = False | |
| ch_axis: int = 0 | |
| group_size: Optional[int] = None # for group-wise along last dim (Linear style) | |
| def _signed_qrange(bits: int) -> Tuple[int, int]: | |
| qmax = (1 << (bits - 1)) - 1 | |
| qmin = - (1 << (bits - 1)) | |
| return qmin, qmax | |
| def calc_qparams_symmetric( | |
| x: torch.Tensor, | |
| bits: int, | |
| per_channel: bool, | |
| ch_axis: int, | |
| clip_percentile: Optional[float] = None, | |
| sample: Optional[torch.Tensor] = None, | |
| eps: float = 1e-8, | |
| ) -> QParams: | |
| """ | |
| Compute symmetric scale such that x/scale is in [qmin,qmax]. | |
| If clip_percentile is provided, use percentile of |x| instead of max(|x|). | |
| For percentile: pass a 1D sample of abs(x) values (optional); otherwise uses abs(x).flatten() which can be heavy. | |
| """ | |
| qmin, qmax = _signed_qrange(bits) | |
| if per_channel: | |
| # reduce all dims except ch_axis | |
| dims = list(range(x.ndim)) | |
| dims.pop(ch_axis) | |
| if clip_percentile is None: | |
| max_abs = x.abs().amax(dim=dims) | |
| else: | |
| # approximate percentile per-channel by using max as fallback (cheap) | |
| # For exact percentile per-channel, provide 'sample' externally. | |
| max_abs = x.abs().amax(dim=dims) | |
| scale = max_abs / float(qmax) | |
| else: | |
| if clip_percentile is None: | |
| max_abs = x.abs().max() | |
| else: | |
| if sample is None: | |
| # approximate: take a smaller random sample | |
| flat = x.abs().flatten() | |
| if flat.numel() > 4096: | |
| idx = torch.randint(0, flat.numel(), (4096,), device=flat.device) | |
| flat = flat[idx] | |
| else: | |
| flat = sample | |
| k = max(0.0, min(1.0, float(clip_percentile))) | |
| max_abs = torch.quantile(flat, k) | |
| scale = max_abs / float(qmax) | |
| scale = torch.clamp(scale, min=eps) | |
| return QParams(scale=scale, qmax=qmax, qmin=qmin, bits=bits, per_channel=per_channel, ch_axis=ch_axis, group_size=None) | |
| def quantize_symmetric(x: torch.Tensor, qparams: QParams) -> torch.Tensor: | |
| """ | |
| Return int8 tensor containing quantized values (still stored as int8 even for int4). | |
| """ | |
| scale = qparams.scale | |
| if qparams.per_channel: | |
| view_shape = [1] * x.ndim | |
| view_shape[qparams.ch_axis] = -1 | |
| scale_b = scale.view(view_shape) | |
| else: | |
| scale_b = scale | |
| q = torch.round(x / scale_b) | |
| q = torch.clamp(q, qparams.qmin, qparams.qmax) | |
| return q.to(torch.int8) | |
| def dequantize_symmetric(q: torch.Tensor, qparams: QParams) -> torch.Tensor: | |
| scale = qparams.scale | |
| if qparams.per_channel: | |
| view_shape = [1] * q.ndim | |
| view_shape[qparams.ch_axis] = -1 | |
| scale_b = scale.view(view_shape) | |
| else: | |
| scale_b = scale | |
| return q.to(torch.float32) * scale_b.to(torch.float32) | |
| # --------------------------- | |
| # Group-wise weight quant (LLM-style, along last dim of weight row) | |
| # --------------------------- | |
| class GroupQParams: | |
| scale: torch.Tensor # [out, n_groups] | |
| qmax: int | |
| qmin: int | |
| bits: int | |
| group_size: int | |
| in_features: int | |
| padded_in: int | |
| def calc_group_qparams_symmetric(W_2d: torch.Tensor, bits: int, group_size: int, eps: float = 1e-8) -> GroupQParams: | |
| """ | |
| W_2d shape [out, in]. Group along in dimension for each out row. | |
| """ | |
| if W_2d.ndim != 2: | |
| raise ValueError("calc_group_qparams_symmetric expects 2D weight") | |
| out, in_features = W_2d.shape | |
| qmin, qmax = _signed_qrange(bits) | |
| padded_in = int(math.ceil(in_features / group_size) * group_size) | |
| if padded_in != in_features: | |
| pad = padded_in - in_features | |
| Wp = torch.cat([W_2d, torch.zeros((out, pad), device=W_2d.device, dtype=W_2d.dtype)], dim=1) | |
| else: | |
| Wp = W_2d | |
| Wg = Wp.view(out, padded_in // group_size, group_size) | |
| max_abs = Wg.abs().amax(dim=2) # [out, n_groups] | |
| scale = torch.clamp(max_abs / float(qmax), min=eps) | |
| return GroupQParams(scale=scale, qmax=qmax, qmin=qmin, bits=bits, group_size=group_size, in_features=in_features, padded_in=padded_in) | |
| def quantize_group_symmetric(W_2d: torch.Tensor, qparams: GroupQParams) -> torch.Tensor: | |
| out, in_features = W_2d.shape | |
| if in_features != qparams.in_features: | |
| raise ValueError("in_features mismatch") | |
| padded_in = qparams.padded_in | |
| if padded_in != in_features: | |
| pad = padded_in - in_features | |
| Wp = torch.cat([W_2d, torch.zeros((out, pad), device=W_2d.device, dtype=W_2d.dtype)], dim=1) | |
| else: | |
| Wp = W_2d | |
| Wg = Wp.view(out, padded_in // qparams.group_size, qparams.group_size) | |
| scale = qparams.scale.unsqueeze(-1) # [out, n_groups, 1] | |
| q = torch.round(Wg / scale) | |
| q = torch.clamp(q, qparams.qmin, qparams.qmax).to(torch.int8) | |
| return q.view(out, padded_in) | |
| def dequantize_group_symmetric(q_int8: torch.Tensor, qparams: GroupQParams) -> torch.Tensor: | |
| out, padded_in = q_int8.shape | |
| if padded_in != qparams.padded_in: | |
| raise ValueError("padded_in mismatch") | |
| qg = q_int8.view(out, padded_in // qparams.group_size, qparams.group_size).to(torch.float32) | |
| scale = qparams.scale.unsqueeze(-1).to(torch.float32) | |
| W = qg * scale | |
| W = W.view(out, padded_in)[:, : qparams.in_features] | |
| return W | |
| # --------------------------- | |
| # Prob/Softmax power-of-two quantization | |
| # --------------------------- | |
| def quantize_prob_power2(p: torch.Tensor, k: int) -> torch.Tensor: | |
| """ | |
| Quantize probabilities p in [0,1] to multiples of 1/2^k. | |
| """ | |
| k = int(k) | |
| base = 1 << k | |
| q = torch.round(p * base) / base | |
| return torch.clamp(q, 0.0, 1.0) | |