Delete utils/quantization.py with huggingface_hub

utils/quantization.py

@@ -1,585 +0,0 @@
-"""quantization.py — W8A8 Quantization via SmoothQuant para CNN-BiGRU.
-
-Implementa quantização weight-only-8-bit + activation-8-bit (W8A8) usando
-a técnica SmoothQuant (Xiao et al., 2023) para migrar a variância das
-ativações para os pesos, reduzindo a perda de precisão.
-
-==============================================================================
-ANÁLISE MATEMÁTICA E LÓGICA — SmoothQuant + W8A8
-==============================================================================
-
-PROBLEMA:
-    Em modelos LLM, as ativações têm outliers em alguns canais que tornam
-    a quantização INT8 difícil. Se quantizarmos diretamente, os outliers
-    saturam os outros canais, levando a grandes erros.
-
-SOLUÇÃO SmoothQuant:
-    Seja Y = X * W, onde X ∈ R^{B×T×d_in} e W ∈ R^{d_in×d_out}.
-
-    1. Para cada canal de entrada i, computa o máximo absoluto:
-        s_i = max|X_i| / max|W_i|  (em batch, suavizado por alpha)
-
-       Ou mais precisamente:
-        s_i = (max|X_i|^alpha) / (max|W_i|^(1-alpha))
-       com alpha ∈ [0, 1] tipicamente 0.5.
-
-    2. Migra a variância: dividir X por s e multiplicar W por s:
-        X' = X / s   (ativações suavizadas — outliers reduzidos)
-        W' = W * s   (pesos absorvem a escala)
-
-       Como Y = X * W = (X/s) * (s*W) = X' * W', a operação matricial
-       é matematicamente equivalente.
-
-    3. Quantiza ambos para INT8 com escala por tensor ou por canal:
-        X_q = round(X' / scale_x) * scale_x
-        W_q = round(W' / scale_w) * scale_w
-
-    4. Y ≈ X_q * W_q (com erro de quantização reduzido)
-
-W8A8:
-    - W: pesos em INT8 (8-bit weights)
-    - A: ativações em INT8 (8-bit activations)
-    - Redução de memória: ~4x (FP32 -> INT8)
-    - Speedup: 2-4x em hardware com suporte INT8 (AMX, AVX512-VNNI)
-
-INTEGRAÇÃO COM CNN-BiGRU:
-    - Aplicado após o treino (post-training quantization)
-    - Aplicável a: Linear (atenção, FFN, classificador), Conv1d
-    - Não aplicado a: Embedding (mantém FP32 para precisão)
-    - Em CPU sem AMX, ainda economiza memória (sem speedup significativo)
-
-==============================================================================
-USO
-==============================================================================
-
-    from cnn_bigru.utils.quantization import (
-        SmoothQuantizer, W8A8Config, quantize_model_w8a8
-    )
-
-    # Após treino:
-    quantizer = SmoothQuantizer(W8A8Config(alpha=0.5, n_calibration_batches=5))
-    quantized_model = quantizer.quantize(model, calibration_dataloader)
-
-Autor: CNN-BiGRU Project
-"""
-from __future__ import annotations
-
-import logging
-import math
-from dataclasses import dataclass, field
-from typing import Dict, List, Optional, Tuple, Union
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-logger = logging.getLogger(__name__)
-
-
-# ============================================================================
-# Configuração
-# ============================================================================
-
-@dataclass
-class W8A8Config:
-    """Configuração da quantização W8A8 com SmoothQuant."""
-    # Alpha de suavização (0=só pesos, 1=só ativações, 0.5=balanceado)
-    alpha: float = 0.5
-    # Número de batches para calibração
-    n_calibration_batches: int = 5
-    # Tipo de escala: "per_tensor" ou "per_channel"
-    scale_type: str = "per_channel"
-    # Manter权重 em FP32 para quantização dinâmica (default: False = INT8 estático)
-    dynamic_weight: bool = False
-    # Quantizar embeddings (default: False)
-    quantize_embeddings: bool = False
-    # Quantizar Conv1d (default: True)
-    quantize_conv1d: bool = True
-    # Quantizar LayerNorm (default: False — sensível)
-    quantize_layernorm: bool = False
-    # Clip threshold para outliers (em desvios-padrão; None = sem clip)
-    outlier_clip_std: Optional[float] = 4.0
-    # Device para calibração
-    device: str = "cpu"
-    # Simular quantização (mantém FP32 mas aplica ruído de quantização)
-    simulate: bool = False
-
-
-# ============================================================================
-# SmoothQuant Calibrator
-# ============================================================================
-
-class SmoothQuantCalibrator:
-    """Coleta estatísticas de ativações e pesos para SmoothQuant.
-
-    Per-corre o modelo com dados de calibração e registra max|X| e max|W|
-    para cada camada Linear/Conv1d.
-    """
-
-    def __init__(self, config: W8A8Config):
-        self.config = config
-        self.stats: Dict[str, Dict[str, torch.Tensor]] = {}
-
-    def _hook_factory(self, name: str):
-        """Cria um hook forward para coletar max|X|."""
-        def hook(module, input, output):
-            # input é uma tupla; input[0] é o tensor principal
-            if not isinstance(input, tuple) or len(input) == 0:
-                return
-            x = input[0]
-            if not isinstance(x, torch.Tensor):
-                return
-            # Reduzir para [d_in] (assumindo última dim = features)
-            with torch.no_grad():
-                if x.dim() >= 2:
-                    # max sobre todas as dims exceto a última
-                    x_flat = x.reshape(-1, x.size(-1))
-                    max_x = x_flat.abs().max(dim=0).values  # [d_in]
-                else:
-                    max_x = x.abs()  # [d_in]
-                # Acumular max (não é média — pegamos o máximo global)
-                if name not in self.stats:
-                    self.stats[name] = {"max_x": max_x.clone()}
-                else:
-                    self.stats[name]["max_x"] = torch.maximum(
-                        self.stats[name]["max_x"], max_x
-                    )
-        return hook
-
-    def calibrate(
-        self,
-        model: nn.Module,
-        dataloader,
-        n_batches: Optional[int] = None,
-    ) -> Dict[str, Dict[str, torch.Tensor]]:
-        """Coleta estatísticas via forward hooks.
-
-        Args:
-            model: modelo a calibrar
-            dataloader: iterador de batches
-            n_batches: número de batches (default: config.n_calibration_batches)
-
-        Returns:
-            dict {layer_name: {"max_x": [d_in], "max_w": [d_out]}}
-        """
-        n_batches = n_batches or self.config.n_calibration_batches
-        device = self.config.device
-
-        # Registrar hooks em todas as camadas Linear e Conv1d
-        hooks = []
-        layer_modules = []
-        for name, module in model.named_modules():
-            if isinstance(module, (nn.Linear, nn.Conv1d)):
-                # Skip embeddings
-                if "embedding" in name.lower() and not self.config.quantize_embeddings:
-                    continue
-                if isinstance(module, nn.Conv1d) and not self.config.quantize_conv1d:
-                    continue
-                hook = module.register_forward_hook(self._hook_factory(name))
-                hooks.append(hook)
-                layer_modules.append((name, module))
-
-        # Forward pass em modo eval (sem gradientes)
-        model.eval()
-        model.to(device)
-        was_training = model.training
-        model.eval()
-
-        try:
-            with torch.no_grad():
-                for i, batch in enumerate(dataloader):
-                    if i >= n_batches:
-                        break
-                    try:
-                        # Tentar diferentes formatos de batch
-                        if isinstance(batch, dict):
-                            input_ids_a = batch.get("input_ids_a")
-                            input_ids_b = batch.get("input_ids_b")
-                            images = batch.get("images")
-                            audios = batch.get("audios")
-                            if input_ids_a is not None and input_ids_b is not None:
-                                # Tentar forward do modelo multimodal
-                                try:
-                                    model(
-                                        input_ids_a.to(device),
-                                        input_ids_b.to(device),
-                                        images=images.to(device) if images is not None else None,
-                                        audios=audios.to(device) if audios is not None else None,
-                                        mode="classify",
-                                    )
-                                except Exception:
-                                    # Fallback: forward sem imagens/áudios
-                                    model(input_ids_a.to(device), input_ids_b.to(device))
-                        elif isinstance(batch, (list, tuple)) and len(batch) >= 2:
-                            model(batch[0].to(device), batch[1].to(device))
-                        else:
-                            logger.debug(f"Batch format não reconhecido: {type(batch)}")
-                    except Exception as e:
-                        logger.debug(f"Calibração batch {i} falhou: {e}")
-                        continue
-        finally:
-            # Remover hooks
-            for hook in hooks:
-                hook.remove()
-            if was_training:
-                model.train()
-
-        # Coletar max|W| para cada camada
-        for name, module in layer_modules:
-            if name not in self.stats:
-                continue
-            w = module.weight.data
-            if isinstance(module, nn.Linear):
-                # w: [d_out, d_in] -> max sobre d_out
-                max_w = w.abs().amax(dim=0)  # [d_in]
-            elif isinstance(module, nn.Conv1d):
-                # w: [out_channels, in_channels//groups, kernel_size]
-                # max sobre out_channels e kernel_size
-                max_w = w.abs().amax(dim=(0, 2))  # [in_channels]
-            else:
-                continue
-            self.stats[name]["max_w"] = max_w.clone()
-
-        logger.info(
-            "SmoothQuant calibrado: %d camadas, %d batches",
-            len(self.stats), n_batches,
-        )
-        return self.stats
-
-
-# ============================================================================
-# SmoothQuantizer
-# ============================================================================
-
-class SmoothQuantizer:
-    """Aplica quantização W8A8 com SmoothQuant a um modelo.
-
-    Args:
-        config: configuração W8A8
-    """
-
-    def __init__(self, config: W8A8Config):
-        self.config = config
-        self.calibrator = SmoothQuantCalibrator(config)
-
-    def compute_scales(
-        self,
-        stats: Dict[str, Dict[str, torch.Tensor]],
-    ) -> Dict[str, torch.Tensor]:
-        """Computa fatores de escala s_i = (max_x^alpha) / (max_w^(1-alpha)).
-
-        Args:
-            stats: dict {layer_name: {"max_x": [d_in], "max_w": [d_in]}}
-
-        Returns:
-            dict {layer_name: scale [d_in]}
-        """
-        scales = {}
-        alpha = self.config.alpha
-        eps = 1e-8
-
-        for name, s in stats.items():
-            max_x = s["max_x"].float()
-            max_w = s.get("max_w")
-            if max_w is None:
-                # Sem max_w (não calibrado), usa só max_x
-                scale = torch.ones_like(max_x)
-            else:
-                max_w = max_w.float().to(max_x.device)
-                # s = (max_x^alpha) / (max_w^(1-alpha))
-                # Adiciona eps para evitar divisão por zero
-                scale = (max_x.clamp(min=eps).pow(alpha) /
-                         max_w.clamp(min=eps).pow(1 - alpha))
-                # Clipa outliers se configurado
-                if self.config.outlier_clip_std is not None:
-                    mean = scale.mean()
-                    std = scale.std()
-                    threshold = self.config.outlier_clip_std * std
-                    scale = scale.clamp(min=mean - threshold, max=mean + threshold)
-                # Normaliza para média 1 (preserva escala global)
-                scale = scale * (scale.numel() / scale.sum().clamp(min=eps))
-            scales[name] = scale
-
-        return scales
-
-    def quantize_tensor_per_channel(
-        self,
-        tensor: torch.Tensor,
-        scale: Optional[torch.Tensor] = None,
-        n_bits: int = 8,
-        axis: int = -1,
-    ) -> torch.Tensor:
-        """Quantiza um tensor para INT8 (per-channel ou per-tensor).
-
-        Args:
-            tensor: tensor a quantizar
-            scale: escala por canal (se None, calcula automaticamente)
-                Se fornecido, deve ter shape broadcastable com `tensor`.
-                Para per-row em [d_out, d_in], use scale shape [d_out, 1].
-            n_bits: número de bits (default 8)
-            axis: eixo para per-channel (apenas quando scale=None e é 1D)
-
-        Returns:
-            tensor quantizado (dequantizado para FP32 para uso em forward)
-        """
-        qmax = 2 ** (n_bits - 1) - 1  # 127 para INT8 simétrico
-        qmin = -qmax
-
-        if scale is None:
-            # Per-tensor
-            max_abs = tensor.abs().max()
-            scale = max_abs.clamp(min=1e-8) / qmax
-            # Quantize-dequantize
-            q = torch.round(tensor / scale).clamp(qmin, qmax)
-            return q * scale
-        else:
-            # Per-channel: scale deve ser broadcastable com tensor
-            scale = scale.to(tensor.device)
-            # Se scale é 1D, expandimos para o eixo
-            if scale.dim() == 1:
-                shape = [1] * tensor.dim()
-                shape[axis] = scale.size(0)
-                scale_b = scale.view(shape)
-            else:
-                # scale já tem shape broadcastable
-                scale_b = scale
-            q = torch.round(tensor / scale_b).clamp(qmin, qmax)
-            return q * scale_b
-
-    def quantize_model(
-        self,
-        model: nn.Module,
-        dataloader=None,
-    ) -> nn.Module:
-        """Aplica quantização W8A8 ao modelo.
-
-        Args:
-            model: modelo a quantizar
-            dataloader: dados de calibração (necessário para SmoothQuant estático)
-
-        Returns:
-            modelo quantizado (parâmetros substituídos por versões INT8 simuladas)
-        """
-        config = self.config
-        device = config.device
-        model.to(device)
-
-        # 1. Calibrar se dataloader fornecido
-        if dataloader is not None:
-            stats = self.calibrator.calibrate(model, dataloader)
-            scales = self.compute_scales(stats)
-        else:
-            stats = {}
-            scales = {}
-
-        # 2. Aplicar SmoothQuant + quantização W8A8 a cada camada
-        n_quantized = 0
-        n_skipped = 0
-        with torch.no_grad():
-            for name, module in model.named_modules():
-                if not isinstance(module, (nn.Linear, nn.Conv1d)):
-                    continue
-                if "embedding" in name.lower() and not config.quantize_embeddings:
-                    n_skipped += 1
-                    continue
-                if isinstance(module, nn.Conv1d) and not config.quantize_conv1d:
-                    n_skipped += 1
-                    continue
-
-                # Aplicar SmoothQuant: W' = W * s (migrar escala dos pesos)
-                scale = scales.get(name)
-                w = module.weight.data
-                if scale is not None:
-                    # Suavizar pesos: multiplicar pela escala
-                    if isinstance(module, nn.Linear):
-                        # w: [d_out, d_in], scale: [d_in]
-                        w_smoothed = w * scale.to(w.device).unsqueeze(0)
-                    elif isinstance(module, nn.Conv1d):
-                        # w: [out_ch, in_ch//groups, kernel_size], scale: [in_ch]
-                        w_smoothed = w * scale.to(w.device).view(1, -1, 1)
-                    else:
-                        w_smoothed = w
-                else:
-                    w_smoothed = w
-
-                # Quantizar pesos para INT8 (simulado — dequantiza de volta)
-                if config.scale_type == "per_channel":
-                    if isinstance(module, nn.Linear):
-                        # Per-output-channel para pesos
-                        w_scale = w_smoothed.abs().amax(dim=1) / 127.0
-                        w_scale = w_scale.clamp(min=1e-8)
-                        w_q = self.quantize_tensor_per_channel(
-                            w_smoothed, scale=w_scale.unsqueeze(1), axis=1,
-                        )
-                    elif isinstance(module, nn.Conv1d):
-                        w_scale = w_smoothed.abs().amax(dim=(1, 2)) / 127.0
-                        w_scale = w_scale.clamp(min=1e-8)
-                        w_q = self.quantize_tensor_per_channel(
-                            w_smoothed, scale=w_scale.view(-1, 1, 1), axis=0,
-                        )
-                else:
-                    # Per-tensor
-                    w_q = self.quantize_tensor_per_channel(w_smoothed)
-
-                module.weight.data = w_q.to(module.weight.dtype)
-                n_quantized += 1
-
-        logger.info(
-            "SmoothQuant W8A8 aplicado: %d camadas quantizadas, %d ignoradas",
-            n_quantized, n_skipped,
-        )
-
-        # Marcar modelo como quantizado (para uso futuro)
-        model._is_quantized_w8a8 = True  # type: ignore
-        model._quantization_config = config  # type: ignore
-
-        return model
-
-    @staticmethod
-    def is_quantized(model: nn.Module) -> bool:
-        """Verifica se um modelo foi quantizado."""
-        return getattr(model, "_is_quantized_w8a8", False)
-
-
-# ============================================================================
-# Helper: quantize modelo inteiro
-# ============================================================================
-
-def quantize_model_w8a8(
-    model: nn.Module,
-    dataloader=None,
-    alpha: float = 0.5,
-    n_calibration_batches: int = 5,
-    device: str = "cpu",
-    **kwargs,
-) -> nn.Module:
-    """Atalho para quantizar um modelo W8A8.
-
-    Args:
-        model: modelo a quantizar
-        dataloader: dados de calibração (None para quantização dinâmica)
-        alpha: fator SmoothQuant (0.5 default)
-        n_calibration_batches: batches para calibração
-        device: device para calibração
-        **kwargs: outros parâmetros de W8A8Config
-
-    Returns:
-        modelo quantizado
-    """
-    config = W8A8Config(
-        alpha=alpha,
-        n_calibration_batches=n_calibration_batches,
-        device=device,
-        **kwargs,
-    )
-    quantizer = SmoothQuantizer(config)
-    return quantizer.quantize_model(model, dataloader=dataloader)
-
-
-# ============================================================================
-# Helper: estimar redução de memória
-# ============================================================================
-
-def estimate_memory_savings(model: nn.Module) -> Dict[str, float]:
-    """Estima redução de memória após quantização W8A8.
-
-    Como a quantização é SIMULADA (dequantiza de volta para FP32 para uso em
-    forward), a memória real não muda. Esta função reporta a redução POTENCIAL
-    se os pesos fossem armazenados como INT8 real.
-
-    Args:
-        model: modelo (preferencialmente quantizado)
-
-    Returns:
-        dict com tamanhos em MB
-    """
-    is_quantized = getattr(model, "_is_quantized_w8a8", False)
-    fp32_bytes = 0  # embeddings + biases (sempre FP32)
-    quantizable_bytes = 0  # pesos que seriam INT8
-
-    for name, param in model.named_parameters():
-        n = param.numel()
-        if "embedding" in name.lower():
-            # Embeddings mantêm FP32
-            fp32_bytes += n * 4
-        elif "bias" in name.lower():
-            # Biases mantêm FP32 (típico em W8A8)
-            fp32_bytes += n * 4
-        elif param.dim() >= 2:
-            # Pesos de matrizes (Linear, Conv) — quantizáveis
-            quantizable_bytes += n
-        else:
-            # Outros 1D — mantêm FP32
-            fp32_bytes += n * 4
-
-    # Se quantizado: pesos seriam INT8 (1 byte cada)
-    # Se não quantizado: pesos seriam FP32 (4 bytes cada)
-    if is_quantized:
-        int8_bytes = quantizable_bytes * 1  # INT8
-    else:
-        int8_bytes = quantizable_bytes * 4  # FP32
-
-    fp32_mb = fp32_bytes / (1024 ** 2)
-    int8_mb = int8_bytes / (1024 ** 2)
-    total_mb = fp32_mb + int8_mb
-    # Sem quantização, tudo seria FP32
-    no_quant_mb = (fp32_bytes + quantizable_bytes * 4) / (1024 ** 2)
-    reduction = (1 - total_mb / no_quant_mb) * 100 if no_quant_mb > 0 else 0
-
-    return {
-        "fp32_mb": fp32_mb,
-        "int8_mb": int8_mb,
-        "total_mb": total_mb,
-        "no_quant_mb": no_quant_mb,
-        "reduction_pct": reduction,
-        "is_quantized": is_quantized,
-    }
-
-
-# ============================================================================
-# Self-test
-# ============================================================================
-
-def _self_test():
-    """Teste rápido da quantização W8A8."""
-    torch.manual_seed(42)
-
-    # Modelo simples para teste
-    class TinyModel(nn.Module):
-        def __init__(self):
-            super().__init__()
-            self.fc1 = nn.Linear(32, 64)
-            self.fc2 = nn.Linear(64, 32)
-            self.embedding = nn.Embedding(100, 32)
-
-        def forward(self, x):
-            return self.fc2(torch.relu(self.fc1(x)))
-
-    model = TinyModel()
-    print(f"Antes: {sum(p.numel() for p in model.parameters())} params")
-
-    # Quantizar sem dataloader (dinâmico)
-    quantized = quantize_model_w8a8(model, dataloader=None, alpha=0.5)
-    print(f"Quantizado: {SmoothQuantizer.is_quantized(quantized)}")
-
-    # Verificar forward ainda funciona
-    x = torch.randn(2, 32)
-    out = quantized(x)
-    print(f"Output shape: {out.shape}")
-
-    # Estimar savings
-    savings = estimate_memory_savings(quantized)
-    print(f"Memory savings: {savings}")
-
-
-if __name__ == "__main__":
-    _self_test()
-
-
-__all__ = [
-    "W8A8Config",
-    "SmoothQuantCalibrator",
-    "SmoothQuantizer",
-    "quantize_model_w8a8",
-    "estimate_memory_savings",
-]