multimodalart's picture
multimodalart HF Staff
Upload folder using huggingface_hub
fb5b18e verified
Raw History Blame Contribute Delete
21.9 kB
# quantization/cfrn/quantlib/qmodules.py
from __future__ import annotations
import math
from typing import Optional, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F
from .quant_utils import (
QParams,
GroupQParams,
calc_qparams_symmetric,
calc_group_qparams_symmetric,
quantize_symmetric,
dequantize_symmetric,
quantize_group_symmetric,
dequantize_group_symmetric,
pack_int4,
unpack_int4,
)
def _get_compute_dtype(name: str) -> torch.dtype:
name = str(name).lower()
if name in ("fp16", "float16", "half"):
return torch.float16
return torch.float32
class FP32LayerNorm(nn.LayerNorm):
def forward(self, x: torch.Tensor) -> torch.Tensor:
y = F.layer_norm(
x.float(),
self.normalized_shape,
self.weight.float() if self.weight is not None else None,
self.bias.float() if self.bias is not None else None,
self.eps,
)
return y.to(x.dtype)
class FP32GroupNorm(nn.GroupNorm):
def forward(self, x: torch.Tensor) -> torch.Tensor:
y = F.group_norm(
x.float(),
self.num_groups,
self.weight.float() if self.weight is not None else None,
self.bias.float() if self.bias is not None else None,
self.eps,
)
return y.to(x.dtype)
# -------------------------
# Helpers
# -------------------------
def _pad_last_dim(x: torch.Tensor, multiple: int) -> Tuple[torch.Tensor, int]:
"""Pad last dim to be divisible by `multiple`, return (padded, original_C)."""
C = x.shape[-1]
if multiple <= 0 or C % multiple == 0:
return x, C
pad = multiple - (C % multiple)
return F.pad(x, (0, pad)), C
def _kthvalue_quantile(x: torch.Tensor, percentile: float, dim: int, keepdim: bool = True) -> torch.Tensor:
"""
IMPORTANT: keep original semantics: k = ceil(p*n), kthvalue uses 1-indexed k.
For group_size=32, p=0.99 -> k=32 -> max (effectively no clip).
"""
p = float(percentile)
if p >= 1.0:
return x.amax(dim=dim, keepdim=keepdim)
if p <= 0.0:
return x.amin(dim=dim, keepdim=keepdim)
n = x.shape[dim]
k = int(math.ceil(p * n))
k = max(1, min(n, k))
return torch.kthvalue(x, k=k, dim=dim, keepdim=keepdim).values
def _dstg_clip_value(x: torch.Tensor, dstg_clip_percentile: Optional[float], dim: int = -1) -> torch.Tensor:
"""
DSTG clip:
- default: max(|x|)
- optional: kthvalue percentile on |x|
(RMS cap 已删除)
"""
abs_x = x.abs()
if dstg_clip_percentile is None:
clip = abs_x.amax(dim=dim, keepdim=True)
else:
clip = _kthvalue_quantile(abs_x, float(dstg_clip_percentile), dim=dim, keepdim=True)
return torch.clamp(clip, min=1e-8)
def _token_clip_value(x: torch.Tensor, token_clip_percentile: Optional[float], dim: int = -1) -> torch.Tensor:
"""Deprecated alias of `_dstg_clip_value` for compatibility."""
return _dstg_clip_value(x, token_clip_percentile, dim=dim)
def _apply_dcrp_constraints(
clip: torch.Tensor,
x: torch.Tensor,
qmax: int,
tau: Optional[float],
zero_ratio: Optional[float],
dim: int = -1,
) -> torch.Tensor:
"""
DCRP = joint constraint (不拆):
- tau: step/std <= tau => clip <= qmax * tau * std
- zero_ratio: approx P(q=0) <= zr => clip <= 2*qmax*Q_zr(|x|)
只有当 tau 和 zero_ratio 都提供时才启用(避免“拆开”单独生效)。
"""
if tau is None or zero_ratio is None:
return torch.clamp(clip, min=1e-8)
std = x.float().std(dim=dim, keepdim=True, unbiased=False) + 1e-12
c_tau = float(qmax) * float(tau) * std.to(clip.dtype)
zr = float(zero_ratio)
zr = max(0.0, min(1.0, zr))
thr = _kthvalue_quantile(x.abs(), zr, dim=dim, keepdim=True)
c_zr = 2.0 * float(qmax) * thr.to(clip.dtype)
out = torch.minimum(clip, c_tau)
out = torch.minimum(out, c_zr)
return torch.clamp(out, min=1e-8)
def _apply_star_constraints(
clip: torch.Tensor,
x: torch.Tensor,
qmax: int,
tau: Optional[float],
zero_ratio: Optional[float],
dim: int = -1,
) -> torch.Tensor:
"""Deprecated alias of `_apply_dcrp_constraints` for compatibility."""
return _apply_dcrp_constraints(clip=clip, x=x, qmax=qmax, tau=tau, zero_ratio=zero_ratio, dim=dim)
def _keep_topk_by_abs(res: torch.Tensor, k: int, dim: int = -1) -> torch.Tensor:
"""
EP residual topk:
- k<=0: keep all
- k>0: only keep top-k by |res| along `dim`
"""
if k is None or int(k) <= 0:
return res
k = int(k)
n = res.shape[dim]
if k >= n:
return res
idx = res.abs().topk(k=k, dim=dim, largest=True).indices
out = torch.zeros_like(res)
out.scatter_(dim, idx, res.gather(dim, idx))
return out
# -------------------------
# QuantLinear (core: DSTG + DCRP + optional EP)
# -------------------------
class QuantLinear(nn.Module):
def __init__(
self,
linear: nn.Linear,
w_bits: int = 4,
a_bits: int = 4,
# weight quant
w_per_channel: bool = True,
w_group_size: Optional[int] = None,
w_clip_percentile: Optional[float] = None,
# activation quant
act_mode: str = "per_dstg_group", # per_tensor / per_channel / per_dstg_group
act_group_size: int = 32, # DSTG group size
dstg_group_size: Optional[int] = None,
act_per_channel: bool = False, # for per_tensor/per_channel fallback
a_clip_percentile: Optional[float] = None,
# DSTG clip
dstg_clip_percentile: Optional[float] = None,
# Deprecated alias
token_clip_percentile: Optional[float] = None,
# EP
token_outlier_bypass: bool = False,
token_outlier_topk: int = 0, # 0=keep full residual; 4=engineer-friendly
# DCRP (joint)
dcrp_tau: Optional[float] = None,
dcrp_zero_ratio: Optional[float] = None,
# Deprecated aliases
token_step_over_std: Optional[float] = None, # tau
token_zero_ratio: Optional[float] = None, # zr
compute_dtype: str = "fp16",
full_name: Optional[str] = None,
):
super().__init__()
if not isinstance(linear, nn.Linear):
raise TypeError("QuantLinear expects an nn.Linear")
self.full_name = str(full_name) if full_name is not None else ""
self.in_features = int(linear.in_features)
self.out_features = int(linear.out_features)
self.w_bits = int(w_bits)
self.a_bits = int(a_bits)
self.w_per_channel = bool(w_per_channel)
self.w_group_size = int(w_group_size) if w_group_size is not None else None
self.w_clip_percentile = w_clip_percentile
mode = str(act_mode).strip().lower()
if mode in ("per_token_group", "tokengroup", "token_group", "tgq", "dstg"):
mode = "per_dstg_group"
self.act_mode = mode
if dstg_group_size is not None:
act_group_size = int(dstg_group_size)
self.act_group_size = int(act_group_size)
self.dstg_group_size = int(self.act_group_size)
self.act_per_channel = bool(act_per_channel)
self.a_clip_percentile = a_clip_percentile
if dstg_clip_percentile is None and token_clip_percentile is not None:
dstg_clip_percentile = token_clip_percentile
if token_clip_percentile is None and dstg_clip_percentile is not None:
token_clip_percentile = dstg_clip_percentile
self.dstg_clip_percentile = dstg_clip_percentile
self.token_clip_percentile = token_clip_percentile
self.token_outlier_bypass = bool(token_outlier_bypass)
self.token_outlier_topk = int(token_outlier_topk)
if dcrp_tau is None and token_step_over_std is not None:
dcrp_tau = token_step_over_std
if dcrp_zero_ratio is None and token_zero_ratio is not None:
dcrp_zero_ratio = token_zero_ratio
self.dcrp_tau = dcrp_tau
self.dcrp_zero_ratio = dcrp_zero_ratio
self.token_step_over_std = dcrp_tau
self.token_zero_ratio = dcrp_zero_ratio
self.compute_dtype = str(compute_dtype)
# bias (fp32)
if linear.bias is None:
self.register_buffer("bias_fp", None)
else:
self.register_buffer("bias_fp", linear.bias.detach().to(torch.float32))
# -------- weight quant (static) --------
W = linear.weight.detach().to(torch.float32)
if self.w_group_size is not None:
gq = calc_group_qparams_symmetric(W, bits=self.w_bits, group_size=self.w_group_size)
q_int8 = quantize_group_symmetric(W, gq)
self.register_buffer("w_scale_group", gq.scale.to(torch.float32))
self.register_buffer(
"w_group_meta",
torch.tensor([gq.group_size, gq.in_features, gq.padded_in], dtype=torch.int32, device=W.device),
)
if self.w_bits == 4:
packed, shape = pack_int4(q_int8.to(torch.int8))
self.register_buffer("w_q_packed", packed)
self.register_buffer("w_q_shape", torch.tensor(list(shape), dtype=torch.int32, device=W.device))
self.register_buffer("w_q_int8", None)
else:
self.register_buffer("w_q_int8", q_int8.to(torch.int8))
self.register_buffer("w_q_packed", None)
self.register_buffer("w_q_shape", None)
else:
qparams = calc_qparams_symmetric(
W,
bits=self.w_bits,
per_channel=self.w_per_channel,
ch_axis=0,
clip_percentile=self.w_clip_percentile,
)
q_int8 = quantize_symmetric(W, qparams).to(torch.int8)
self.register_buffer("w_scale", qparams.scale.to(torch.float32))
self.register_buffer(
"w_qparams_meta",
torch.tensor([qparams.qmin, qparams.qmax, int(self.w_per_channel)], dtype=torch.int32, device=W.device),
)
if self.w_bits == 4:
packed, shape = pack_int4(q_int8)
self.register_buffer("w_q_packed", packed)
self.register_buffer("w_q_shape", torch.tensor(list(shape), dtype=torch.int32, device=W.device))
self.register_buffer("w_q_int8", None)
else:
self.register_buffer("w_q_int8", q_int8)
self.register_buffer("w_q_packed", None)
self.register_buffer("w_q_shape", None)
def _dequant_weight(self) -> torch.Tensor:
if self.w_group_size is not None:
meta = self.w_group_meta
group_size = int(meta[0].item())
in_features = int(meta[1].item())
padded_in = int(meta[2].item())
if self.w_bits == 4:
shape = tuple(int(x.item()) for x in self.w_q_shape)
q = unpack_int4(self.w_q_packed, shape).to(torch.int8)
else:
q = self.w_q_int8
gq = GroupQParams(
scale=self.w_scale_group,
qmax=(1 << (self.w_bits - 1)) - 1,
qmin=-(1 << (self.w_bits - 1)),
bits=self.w_bits,
group_size=group_size,
in_features=in_features,
padded_in=padded_in,
)
return dequantize_group_symmetric(q, gq)
qmin, qmax, per_ch = [int(x.item()) for x in self.w_qparams_meta]
per_channel = bool(per_ch)
if self.w_bits == 4:
shape = tuple(int(x.item()) for x in self.w_q_shape)
q = unpack_int4(self.w_q_packed, shape).to(torch.int8)
else:
q = self.w_q_int8
qparams = QParams(
scale=self.w_scale,
qmax=qmax,
qmin=qmin,
bits=self.w_bits,
per_channel=per_channel,
ch_axis=0,
)
return dequantize_symmetric(q, qparams)
def _quant_dequant_act_per_tensor(self, x: torch.Tensor) -> torch.Tensor:
if self.a_bits >= 16:
return x
qparams = calc_qparams_symmetric(
x,
bits=self.a_bits,
per_channel=self.act_per_channel,
ch_axis=-1,
clip_percentile=self.a_clip_percentile,
)
q = quantize_symmetric(x, qparams)
return dequantize_symmetric(q, qparams).to(x.dtype)
def _quant_dequant_act_dstg(self, x: torch.Tensor) -> torch.Tensor:
if self.a_bits >= 16:
return x
if x.ndim < 2:
return self._quant_dequant_act_per_tensor(x)
qmin = -(1 << (self.a_bits - 1))
qmax = (1 << (self.a_bits - 1)) - 1
g = max(1, int(self.act_group_size))
x_pad, C0 = _pad_last_dim(x, g)
xg = x_pad.view(*x_pad.shape[:-1], x_pad.shape[-1] // g, g) # [..., nG, g]
clip = _dstg_clip_value(xg, self.dstg_clip_percentile, dim=-1)
clip = _apply_dcrp_constraints(
clip=clip,
x=xg,
qmax=qmax,
tau=self.dcrp_tau,
zero_ratio=self.dcrp_zero_ratio,
dim=-1,
)
x_clip = xg.clamp(min=-clip, max=clip)
if self.token_outlier_bypass:
x_res = xg - x_clip
if self.token_outlier_topk > 0:
x_res = _keep_topk_by_abs(x_res, k=min(self.token_outlier_topk, x_res.shape[-1]), dim=-1)
else:
x_res = 0.0
scale = torch.clamp(clip / float(qmax), min=1e-8)
q = torch.round(x_clip / scale).clamp(qmin, qmax).to(torch.int8)
x_hat = (q.to(torch.float32) * scale.to(torch.float32)).to(xg.dtype)
y = (x_hat + x_res).view(*x_pad.shape)[..., :C0]
return y
def _quant_dequant_act(self, x: torch.Tensor) -> torch.Tensor:
mode = str(self.act_mode).lower()
if mode in ("per_dstg_group", "per_token_group", "tokengroup", "tgq"):
return self._quant_dequant_act_dstg(x)
return self._quant_dequant_act_per_tensor(x)
def forward(self, x: torch.Tensor) -> torch.Tensor:
xq = self._quant_dequant_act(x)
W = self._dequant_weight()
dtype = _get_compute_dtype(self.compute_dtype)
xq = xq.to(dtype)
W = W.to(dtype)
b = None if self.bias_fp is None else self.bias_fp.to(dtype)
return F.linear(xq, W, b)
# -------------------------
# QuantConv2d / QuantConv1d (simple)
# -------------------------
class QuantConv2d(nn.Module):
def __init__(
self,
conv: nn.Conv2d,
w_bits: int = 4,
a_bits: int = 4,
w_per_channel: bool = True,
w_clip_percentile: Optional[float] = None,
act_per_channel: bool = False,
a_clip_percentile: Optional[float] = None,
compute_dtype: str = "fp16",
full_name: Optional[str] = None,
):
super().__init__()
if not isinstance(conv, nn.Conv2d):
raise TypeError("QuantConv2d expects nn.Conv2d")
self.full_name = str(full_name) if full_name is not None else ""
self.w_bits = int(w_bits)
self.a_bits = int(a_bits)
self.w_per_channel = bool(w_per_channel)
self.w_clip_percentile = w_clip_percentile
self.act_per_channel = bool(act_per_channel)
self.a_clip_percentile = a_clip_percentile
self.compute_dtype = str(compute_dtype)
self.stride = conv.stride
self.padding = conv.padding
self.dilation = conv.dilation
self.groups = conv.groups
if conv.bias is None:
self.register_buffer("bias_fp", None)
else:
self.register_buffer("bias_fp", conv.bias.detach().to(torch.float32))
W = conv.weight.detach().to(torch.float32)
qparams = calc_qparams_symmetric(
W, bits=self.w_bits, per_channel=self.w_per_channel, ch_axis=0, clip_percentile=self.w_clip_percentile
)
q_int8 = quantize_symmetric(W, qparams).to(torch.int8)
self.register_buffer("w_scale", qparams.scale.to(torch.float32))
self.register_buffer(
"w_qparams_meta",
torch.tensor([qparams.qmin, qparams.qmax, int(self.w_per_channel)], dtype=torch.int32, device=W.device),
)
if self.w_bits == 4:
packed, shape = pack_int4(q_int8)
self.register_buffer("w_q_packed", packed)
self.register_buffer("w_q_shape", torch.tensor(list(shape), dtype=torch.int32, device=W.device))
self.register_buffer("w_q_int8", None)
else:
self.register_buffer("w_q_int8", q_int8)
self.register_buffer("w_q_packed", None)
self.register_buffer("w_q_shape", None)
def _dequant_weight(self) -> torch.Tensor:
qmin, qmax, per_ch = [int(x.item()) for x in self.w_qparams_meta]
per_channel = bool(per_ch)
if self.w_bits == 4:
shape = tuple(int(x.item()) for x in self.w_q_shape)
q = unpack_int4(self.w_q_packed, shape).to(torch.int8)
else:
q = self.w_q_int8
qparams = QParams(scale=self.w_scale, qmax=qmax, qmin=qmin, bits=self.w_bits, per_channel=per_channel, ch_axis=0)
return dequantize_symmetric(q, qparams)
def _quant_dequant_act(self, x: torch.Tensor) -> torch.Tensor:
if self.a_bits >= 16:
return x
qparams = calc_qparams_symmetric(
x, bits=self.a_bits, per_channel=self.act_per_channel, ch_axis=1, clip_percentile=self.a_clip_percentile
)
q = quantize_symmetric(x, qparams)
return dequantize_symmetric(q, qparams).to(x.dtype)
def forward(self, x: torch.Tensor) -> torch.Tensor:
xq = self._quant_dequant_act(x)
W = self._dequant_weight()
dtype = _get_compute_dtype(self.compute_dtype)
xq = xq.to(dtype)
W = W.to(dtype)
b = None if self.bias_fp is None else self.bias_fp.to(dtype)
return F.conv2d(xq, W, b, stride=self.stride, padding=self.padding, dilation=self.dilation, groups=self.groups)
class QuantConv1d(nn.Module):
def __init__(
self,
conv: nn.Conv1d,
w_bits: int = 4,
a_bits: int = 4,
w_per_channel: bool = True,
w_clip_percentile: Optional[float] = None,
act_per_channel: bool = False,
a_clip_percentile: Optional[float] = None,
compute_dtype: str = "fp16",
full_name: Optional[str] = None,
):
super().__init__()
if not isinstance(conv, nn.Conv1d):
raise TypeError("QuantConv1d expects nn.Conv1d")
self.full_name = str(full_name) if full_name is not None else ""
self.w_bits = int(w_bits)
self.a_bits = int(a_bits)
self.w_per_channel = bool(w_per_channel)
self.w_clip_percentile = w_clip_percentile
self.act_per_channel = bool(act_per_channel)
self.a_clip_percentile = a_clip_percentile
self.compute_dtype = str(compute_dtype)
self.stride = conv.stride
self.padding = conv.padding
self.dilation = conv.dilation
self.groups = conv.groups
if conv.bias is None:
self.register_buffer("bias_fp", None)
else:
self.register_buffer("bias_fp", conv.bias.detach().to(torch.float32))
W = conv.weight.detach().to(torch.float32)
qparams = calc_qparams_symmetric(
W, bits=self.w_bits, per_channel=self.w_per_channel, ch_axis=0, clip_percentile=self.w_clip_percentile
)
q_int8 = quantize_symmetric(W, qparams).to(torch.int8)
self.register_buffer("w_scale", qparams.scale.to(torch.float32))
self.register_buffer(
"w_qparams_meta",
torch.tensor([qparams.qmin, qparams.qmax, int(self.w_per_channel)], dtype=torch.int32, device=W.device),
)
if self.w_bits == 4:
packed, shape = pack_int4(q_int8)
self.register_buffer("w_q_packed", packed)
self.register_buffer("w_q_shape", torch.tensor(list(shape), dtype=torch.int32, device=W.device))
self.register_buffer("w_q_int8", None)
else:
self.register_buffer("w_q_int8", q_int8)
self.register_buffer("w_q_packed", None)
self.register_buffer("w_q_shape", None)
def _dequant_weight(self) -> torch.Tensor:
qmin, qmax, per_ch = [int(x.item()) for x in self.w_qparams_meta]
per_channel = bool(per_ch)
if self.w_bits == 4:
shape = tuple(int(x.item()) for x in self.w_q_shape)
q = unpack_int4(self.w_q_packed, shape).to(torch.int8)
else:
q = self.w_q_int8
qparams = QParams(scale=self.w_scale, qmax=qmax, qmin=qmin, bits=self.w_bits, per_channel=per_channel, ch_axis=0)
return dequantize_symmetric(q, qparams)
def _quant_dequant_act(self, x: torch.Tensor) -> torch.Tensor:
if self.a_bits >= 16:
return x
qparams = calc_qparams_symmetric(
x, bits=self.a_bits, per_channel=self.act_per_channel, ch_axis=1, clip_percentile=self.a_clip_percentile
)
q = quantize_symmetric(x, qparams)
return dequantize_symmetric(q, qparams).to(x.dtype)
def forward(self, x: torch.Tensor) -> torch.Tensor:
xq = self._quant_dequant_act(x)
W = self._dequant_weight()
dtype = _get_compute_dtype(self.compute_dtype)
xq = xq.to(dtype)
W = W.to(dtype)
b = None if self.bias_fp is None else self.bias_fp.to(dtype)
return F.conv1d(xq, W, b, stride=self.stride, padding=self.padding, dilation=self.dilation, groups=self.groups)