File size: 3,562 Bytes
c2992d0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | """
Xavante - smoothquant_compressor.py
Responsabilidade: Compressao W8A8 via SmoothQuant (Lema 4.1).
Integra a versao validada do FlexNet.
"""
from __future__ import annotations
import logging
import torch
import torch.nn as nn
logger = logging.getLogger(__name__)
class SmoothQuantCompressor(nn.Module):
"""
SmoothQuant (Xiao et al., 2022) migra a escala da ativacao para o peso,
permitindo quantizacao W8A8 sem perda significativa.
W_smooth = W * diag(s)
X_smooth = X * diag(s)^-1
s_j = max|X_j|^alpha / max|W_j|^(1-alpha), alpha=0.5 tipicamente.
"""
def __init__(self, alpha: float = 0.5, n_bits: int = 8, calibration_samples: int = 128):
super().__init__()
self.alpha = alpha
self.n_bits = n_bits
self.calibration_samples = calibration_samples
self.register_buffer("smooth_scale", torch.ones(1))
self._calibrated = False
def calibrate(self, weight: torch.Tensor, activation_samples: torch.Tensor) -> None:
"""
weight: [out, in]
activation_samples: [N, in] (diversas amostras reais)
"""
with torch.no_grad():
# max abs por canal de entrada
w_max = weight.abs().amax(dim=0).clamp(min=1e-5) # [in]
x_max = activation_samples.abs().amax(dim=0).clamp(min=1e-5) # [in]
s = (x_max.pow(self.alpha) / w_max.pow(1 - self.alpha)).clamp(min=1e-5)
self.smooth_scale = s.to(weight.device, weight.dtype)
self._calibrated = True
logger.info("SmoothQuant calibrado: scale mean=%.4f, std=%.4f", s.mean().item(), s.std().item())
def smooth_weight(self, weight: torch.Tensor) -> torch.Tensor:
"""Aplica W_smooth = W * diag(s) (multiplicacao por coluna)."""
if not self._calibrated:
return weight
return weight * self.smooth_scale.unsqueeze(0)
def smooth_activation(self, x: torch.Tensor) -> torch.Tensor:
"""Aplica X_smooth = X / diag(s)."""
if not self._calibrated:
return x
return x / self.smooth_scale
def quantize_per_tensor_symmetric(self, t: torch.Tensor) -> torch.Tensor:
# Quantizacao simetrica per-tensor em n_bits
qmax = 2 ** (self.n_bits - 1) - 1
scale = t.abs().amax().clamp(min=1e-8) / qmax
q = torch.clamp((t / scale).round(), -qmax, qmax)
return q.to(torch.int8) if self.n_bits == 8 else q.to(torch.int16)
def fake_quantize(self, t: torch.Tensor) -> torch.Tensor:
"""Quantizacao fake (diferenciavel via STE) para treino QAT."""
qmax = 2 ** (self.n_bits - 1) - 1
scale = t.abs().amax().clamp(min=1e-8) / qmax
q = torch.clamp((t / scale).round(), -qmax, qmax)
return q * scale # retorna em float para backward
def forward(self, weight: torch.Tensor, x: torch.Tensor, fake: bool = True) -> torch.Tensor:
w_s = self.smooth_weight(weight)
x_s = self.smooth_activation(x)
if fake:
w_q = self.fake_quantize(w_s)
x_q = self.fake_quantize(x_s)
else:
w_q = self.dequantize(self.quantize_per_tensor_symmetric(w_s), w_s)
x_q = self.dequantize(self.quantize_per_tensor_symmetric(x_s), x_s)
return torch.nn.functional.linear(x_q, w_q)
def dequantize(self, qt: torch.Tensor, ref: torch.Tensor) -> torch.Tensor:
qmax = 2 ** (self.n_bits - 1) - 1
scale = ref.abs().amax().clamp(min=1e-8) / qmax
return qt.to(ref.dtype) * scale
__all__ = ["SmoothQuantCompressor"]
|