Spaces:
Running on Zero
Running on Zero
| # quantization/escnet/quant/quantizers.py | |
| from __future__ import annotations | |
| from dataclasses import dataclass | |
| from typing import Optional, Tuple | |
| import torch | |
| from .config import QuantConfig | |
| from .observers import MinMaxObserver | |
| from .dstg_dcrp import dstg_dcrp_qdq | |
| def _qmax(bits: int) -> int: | |
| return (1 << (bits - 1)) - 1 | |
| def symmetric_qdq(x: torch.Tensor, scale: torch.Tensor, qmax: int) -> torch.Tensor: | |
| """Symmetric QDQ: round(x/scale) * scale.""" | |
| orig_dtype = x.dtype | |
| xf = x.float() | |
| sf = scale.float() | |
| # broadcast | |
| q = torch.round(xf / sf).clamp(-qmax, qmax) | |
| return (q * sf).to(orig_dtype) | |
| class StaticActQuant: | |
| """Static activation QDQ (needs calibration).""" | |
| bits: int | |
| mode: str # per_tensor | per_channel | |
| ch_axis: int | |
| observer: MinMaxObserver | |
| scale: Optional[torch.Tensor] = None # on CPU by default | |
| def __post_init__(self) -> None: | |
| if self.mode not in ("per_tensor", "per_channel"): | |
| raise ValueError(f"Unsupported static act mode: {self.mode}") | |
| def observe(self, x: torch.Tensor) -> None: | |
| self.observer.update(x) | |
| def finalize(self, device: torch.device) -> None: | |
| absmax = self.observer.get_absmax(device=device).clamp(min=1e-8) | |
| self.scale = (absmax / float(_qmax(self.bits))).detach() | |
| def __call__(self, x: torch.Tensor) -> torch.Tensor: | |
| if self.scale is None: | |
| # pre-finalize fallback: identity | |
| return x | |
| scale = self.scale.to(x.device) | |
| return symmetric_qdq(x, scale=scale, qmax=_qmax(self.bits)) | |
| class DynamicDSTGDCRPActQuant: | |
| """Dynamic per-DSTG-group + DCRP (no calibration needed).""" | |
| bits: int | |
| group_size: int | |
| dstg_clip_percentile: Optional[float] | |
| dcrp_tau: Optional[float] | |
| dcrp_zero_ratio: Optional[float] | |
| residual_topk: int = 0 | |
| def __call__(self, x: torch.Tensor) -> torch.Tensor: | |
| return dstg_dcrp_qdq( | |
| x, | |
| qmax=_qmax(self.bits), | |
| group_size=self.group_size, | |
| dstg_clip_percentile=self.dstg_clip_percentile, | |
| tau=self.dcrp_tau, | |
| zero_ratio=self.dcrp_zero_ratio, | |
| residual_topk=self.residual_topk, | |
| ) | |
| # Backward compatibility alias | |
| DynamicTGQStarActQuant = DynamicDSTGDCRPActQuant | |
| def quantize_weight_symmetric( | |
| w: torch.Tensor, bits: int, mode: str = "per_channel", ch_axis: int = 0 | |
| ) -> Tuple[torch.Tensor, torch.Tensor]: | |
| """Return (w_int8, w_scale). Symmetric quant, scale is float32.""" | |
| qmax = _qmax(bits) | |
| wf = w.detach().float() | |
| if mode == "per_tensor": | |
| absmax = wf.abs().max().clamp(min=1e-8) | |
| scale = absmax / float(qmax) | |
| wq = torch.round(wf / scale).clamp(-qmax, qmax).to(torch.int8) | |
| return wq, torch.tensor(scale, dtype=torch.float32) | |
| if mode == "per_channel": | |
| # reduce all dims except ch_axis | |
| dims = [d for d in range(wf.dim()) if d != ch_axis] | |
| absmax = wf.abs().amax(dim=dims).clamp(min=1e-8) | |
| scale = absmax / float(qmax) # [Cout] | |
| # reshape for broadcasting | |
| shape = [1] * wf.dim() | |
| shape[ch_axis] = -1 | |
| scale_bc = scale.view(*shape) | |
| wq = torch.round(wf / scale_bc).clamp(-qmax, qmax).to(torch.int8) | |
| return wq, scale.to(torch.float32) | |
| raise ValueError(f"Unsupported weight mode: {mode}") | |
| def make_activation_quantizer(qcfg: QuantConfig, x_example: Optional[torch.Tensor] = None): | |
| """Factory for layer-local activation quantizer.""" | |
| if qcfg.a_mode in ("per_dstg_group", "per_token_group"): | |
| return DynamicDSTGDCRPActQuant( | |
| bits=qcfg.a_bits, | |
| group_size=qcfg.group_size, | |
| dstg_clip_percentile=qcfg.dstg_clip_percentile, | |
| dcrp_tau=qcfg.dcrp_tau, | |
| dcrp_zero_ratio=qcfg.dcrp_zero_ratio, | |
| residual_topk=qcfg.ep_residual_topk, | |
| ) | |
| if qcfg.a_mode == "per_tensor": | |
| return StaticActQuant( | |
| bits=qcfg.a_bits, | |
| mode="per_tensor", | |
| ch_axis=1, | |
| observer=MinMaxObserver(mode="per_tensor", ch_axis=1), | |
| ) | |
| if qcfg.a_mode == "per_channel": | |
| # assume channel axis: for Conv it's dim=1, for Linear inputs it's last dim | |
| ch_axis = 1 | |
| if x_example is not None and x_example.dim() == 2: | |
| ch_axis = -1 | |
| return StaticActQuant( | |
| bits=qcfg.a_bits, | |
| mode="per_channel", | |
| ch_axis=ch_axis, | |
| observer=MinMaxObserver(mode="per_channel", ch_axis=ch_axis), | |
| ) | |
| raise ValueError(f"Unknown a_mode: {qcfg.a_mode}") | |