V6.5: 1000 samples (500+500), MTP val+entropy, EWC+W8A8 eval investigation, integrate gru-ring modules
c2992d0 verified | """ | |
| 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"] | |