File size: 2,543 Bytes
dc9acb9 | 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 | """
HyperSparse Adaptive Quantization — sparsity adaptativa via kthvalue
HSAQ cuantiza ACTIVACIONES a {0, valor} mediante máscara binaria dinámica.
NO es cuantización de pesos. NO es INT8/INT4. NO es bitsandbytes.
Objetivo: optimizar la cuantización mejor que TurboQuant (Google) haciendo
que TODO sea adaptativo, para no usar recursos innecesariamente. Esto permite
usar modelos más grandes que la capacidad del hardware, ejecutando solo los
tokens/neuronas necesarias para el trabajo operativo.
"""
import torch
import torch.nn as nn
class HSAQ(nn.Module):
"""HyperSparse Adaptive Quantization.
Aplica una máscara de sparse adaptativa por batch.
Sparsity escalonada: capas tempranas preservan más información,
capas tardías comprimen más agresivamente.
Args:
sparsity: Fracción de neuronas a enmascarar (0.3 = 30% → 0).
"""
def __init__(self, sparsity: float = 0.3):
super().__init__()
self.sparsity = sparsity
def forward(self, x: torch.Tensor, sparsity_override: float | None = None) -> torch.Tensor:
"""Aplica máscara de sparsity adaptativa por batch.
Args:
x: Tensor de entrada (B, D) o (B, T, D).
sparsity_override: Sparsity para esta llamada (opcional).
Si se pasa, usa este valor en lugar de self.sparsity.
"""
s = sparsity_override if sparsity_override is not None else self.sparsity
if s <= 0.0:
return x
flat = x.abs().view(x.size(0), -1)
n = flat.size(1)
k = max(1, min(n - 1, int(n * s)))
# kthvalue: encuentra el valor en el percentil k (bottom-k%)
# Los valores >= thresh son el top-(1-s)% que pasan
thresh = torch.kthvalue(flat, k, dim=1).values
thresh = thresh.view(-1, *([1] * (x.dim() - 1)))
mask = x.abs() >= thresh
self._last_sparsity = 1.0 - mask.float().mean().item()
self._last_threshold = thresh.view(-1).mean().item()
self._last_sparsity_target = s
return x * mask
def get_stats(self) -> dict:
"""Retorna métricas actuales de HSAQ."""
return {
'sparsity': self.sparsity,
'sparsity_target': getattr(self, '_last_sparsity_target', self.sparsity),
'actual_sparsity': getattr(self, '_last_sparsity', 0.0),
'threshold': getattr(self, '_last_threshold', 0.0),
'type': 'activation_quantization',
'method': 'kthvalue_dynamic_threshold',
}
|