Spaces:
Running on Zero
Running on Zero
Download quantization/cfrn/quantlib/qmodules.py from hugging-apps/cod-tdq-camouflaged-detection: direct link, hf CLI and curl.
- Browser
- Download file 21.9 kB
-
https://huggingface.co/spaces/hugging-apps/cod-tdq-camouflaged-detection/resolve/main/quantization/cfrn/quantlib/qmodules.py
- Command line
-
hf download hf://spaces/hugging-apps/cod-tdq-camouflaged-detection/quantization/cfrn/quantlib/qmodules.py
-
curl -L -o qmodules.py https://huggingface.co/spaces/hugging-apps/cod-tdq-camouflaged-detection/resolve/main/quantization/cfrn/quantlib/qmodules.py
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) | |