| import torch
|
| import torch.nn as nn
|
| import torch.nn.functional as F
|
|
|
| from .modules import sparse as sp
|
|
|
|
|
| def _pack_4bit_tensor(q: torch.Tensor) -> torch.Tensor:
|
| q = q.to(torch.int8)
|
| q = torch.clamp(q, -8, 7)
|
| q_u = ((q + 8) & 0xF).to(torch.uint8).reshape(-1)
|
| if q_u.numel() % 2 == 1:
|
| q_u = torch.cat([q_u, torch.zeros(1, dtype=torch.uint8, device=q_u.device)])
|
| q_u = q_u.view(-1, 2)
|
| return q_u[:, 0] | (q_u[:, 1] << 4)
|
|
|
|
|
| def _unpack_4bit_tensor(packed: torch.Tensor, shape: torch.Size) -> torch.Tensor:
|
| packed = packed.reshape(-1)
|
| low = packed & 0xF
|
| high = packed >> 4
|
| q = torch.empty((packed.numel() * 2,), dtype=torch.int8, device=packed.device)
|
| q[0::2] = low.to(torch.int8)
|
| q[1::2] = high.to(torch.int8)
|
| q = q[: torch.prod(torch.tensor(shape, dtype=torch.int64)).item()]
|
| q = q.view(shape)
|
| return q - 8
|
|
|
|
|
| def _quantize_weight(weight: torch.Tensor, bits: int = 4) -> tuple[torch.Tensor, torch.Tensor]:
|
| if bits != 4:
|
| raise ValueError('Only 4-bit quantization is supported.')
|
|
|
| flat_weight = weight.view(weight.size(0), -1)
|
| max_vals = flat_weight.abs().amax(dim=1)
|
| scales = torch.where(max_vals > 0, max_vals / 7.0, torch.ones_like(max_vals))
|
| scales = scales.to(weight.dtype)
|
|
|
| q_weight = torch.clamp((flat_weight / scales[:, None]).round(), -8, 7).to(torch.int8)
|
| q_weight = q_weight.view(weight.shape)
|
| packed = _pack_4bit_tensor(q_weight)
|
| return packed, scales
|
|
|
|
|
| class QuantizedLinear(nn.Module):
|
| def __init__(
|
| self,
|
| in_features: int,
|
| out_features: int,
|
| bias: bool = True,
|
| dtype: torch.dtype = torch.float16,
|
| bits: int = 4,
|
| ):
|
| super().__init__()
|
| self.in_features = in_features
|
| self.out_features = out_features
|
| self.bits = bits
|
| self.dtype = dtype
|
|
|
| self.register_buffer('weight_packed', torch.empty(0, dtype=torch.uint8))
|
| self.register_buffer('weight_scales', torch.empty((out_features,), dtype=self.dtype))
|
|
|
| if bias:
|
| self.bias = nn.Parameter(torch.zeros(out_features, dtype=self.dtype))
|
| else:
|
| self.bias = None
|
|
|
| @classmethod
|
| def from_float(cls, module: nn.Linear, bits: int = 4, dtype: torch.dtype = torch.float16):
|
| quantized = cls(module.in_features, module.out_features, module.bias is not None, dtype=dtype, bits=bits)
|
| packed, scales = _quantize_weight(module.weight.data, bits=bits)
|
| quantized.weight_packed = packed.to(dtype=torch.uint8)
|
| quantized.weight_scales = scales.to(dtype=dtype)
|
| if module.bias is not None:
|
| quantized.bias.data = module.bias.data.to(dtype=dtype)
|
| return quantized
|
|
|
| def dequantize_weight(self) -> torch.Tensor:
|
| q = _unpack_4bit_tensor(self.weight_packed, torch.Size((self.out_features, self.in_features)))
|
| scales = self.weight_scales.view(self.out_features, *([1] * (q.ndim - 1)))
|
| return q.to(self.dtype) * scales
|
|
|
| def forward(self, input: torch.Tensor) -> torch.Tensor:
|
| input = input.to(self.dtype)
|
| weight = self.dequantize_weight()
|
| return F.linear(input, weight, self.bias)
|
|
|
|
|
| class QuantizedSparseLinear(QuantizedLinear):
|
| def forward(self, input):
|
| feats = input.feats.to(self.dtype)
|
| weight = self.dequantize_weight()
|
| return input.replace(F.linear(feats, weight, self.bias))
|
|
|
|
|
| def quantize_model(model: nn.Module, bits: int = 4, dtype: torch.dtype = torch.float16) -> None:
|
| """Replace linear modules in a model with 4-bit quantized equivalents."""
|
| for name, child in list(model.named_children()):
|
| if isinstance(child, sp.SparseLinear):
|
| model._modules[name] = QuantizedSparseLinear.from_float(child, bits=bits, dtype=dtype)
|
| elif isinstance(child, nn.Linear):
|
| model._modules[name] = QuantizedLinear.from_float(child, bits=bits, dtype=dtype)
|
| else:
|
| quantize_model(child, bits=bits, dtype=dtype)
|
|
|
