import torch import torch.nn as nn try: import quant_cuda except: print('CUDA extension not installed.') def quantize(x, scale, zero, maxq): q = torch.clamp(torch.round(x / scale) + zero, 0, maxq) return scale * (q - zero) class Quantizer(nn.Module): def __init__(self, shape=1): super(Quantizer, self).__init__() self.register_buffer('maxq', torch.tensor(0)) self.register_buffer('scale', torch.zeros(shape)) self.register_buffer('zero', torch.zeros(shape)) def configure( self, bits, perchannel=False, sym=True, mse=False, norm=2.4, grid=100, maxshrink=.8 ): self.maxq = torch.tensor(2 ** bits - 1) self.perchannel = perchannel self.sym = sym self.mse = mse self.norm = norm self.grid = grid self.maxshrink = maxshrink def find_params(self, x, weight=False): dev = x.device self.maxq = self.maxq.to(dev) shape = x.shape if self.perchannel: if weight: x = x.flatten(1) else: if len(shape) == 4: x = x.permute([1, 0, 2, 3]) x = x.flatten(1) if len(shape) == 3: x = x.reshape((-1, shape[-1])).t() if len(shape) == 2: x = x.t() else: x = x.flatten().unsqueeze(0) tmp = torch.zeros(x.shape[0], device=dev) xmin = torch.minimum(x.min(1)[0], tmp) xmax = torch.maximum(x.max(1)[0], tmp) if self.sym: xmax = torch.maximum(torch.abs(xmin), xmax) tmp = xmin < 0 if torch.any(tmp): xmin[tmp] = -xmax[tmp] tmp = (xmin == 0) & (xmax == 0) xmin[tmp] = -1 xmax[tmp] = +1 self.scale = (xmax - xmin) / self.maxq if self.sym: self.zero = torch.full_like(self.scale, (self.maxq + 1) / 2) else: self.zero = torch.round(-xmin / self.scale) if self.mse: best = torch.full([x.shape[0]], float('inf'), device=dev) for i in range(int(self.maxshrink * self.grid)): p = 1 - i / self.grid xmin1 = p * xmin xmax1 = p * xmax scale1 = (xmax1 - xmin1) / self.maxq zero1 = torch.round(-xmin1 / scale1) if not self.sym else self.zero q = quantize(x, scale1.unsqueeze(1), zero1.unsqueeze(1), self.maxq) q -= x q.abs_() q.pow_(self.norm) err = torch.sum(q, 1) tmp = err < best if torch.any(tmp): best[tmp] = err[tmp] self.scale[tmp] = scale1[tmp] self.zero[tmp] = zero1[tmp] if not self.perchannel: if weight: tmp = shape[0] else: tmp = shape[1] if len(shape) != 3 else shape[2] self.scale = self.scale.repeat(tmp) self.zero = self.zero.repeat(tmp) if weight: shape = [-1] + [1] * (len(shape) - 1) # self.scale = self.scale.unsqueeze(1) # self.zero = self.zero.unsqueeze(1) self.scale = self.scale.reshape(shape) self.zero = self.zero.reshape(shape) return if len(shape) == 4: self.scale = self.scale.reshape((1, -1, 1, 1)) self.zero = self.zero.reshape((1, -1, 1, 1)) if len(shape) == 3: self.scale = self.scale.reshape((1, 1, -1)) self.zero = self.zero.reshape((1, 1, -1)) if len(shape) == 2: self.scale = self.scale.unsqueeze(0) self.zero = self.zero.unsqueeze(0) def quantize(self, x): if self.ready(): return quantize(x, self.scale, self.zero, self.maxq) return x def enabled(self): return self.maxq > 0 def ready(self): return torch.all(self.scale != 0) class ActQuantWrapper(nn.Module): def __init__(self, module): super(ActQuantWrapper, self).__init__() self.module = module shape = [1] * len(self.module.weight.shape) if len(shape) == 4: shape[1] = self.module.weight.shape[1] if len(shape) == 3: shape[2] = self.module.weight.shape[2] if len(shape) == 2: shape[1] = self.module.weight.shape[1] self.quantizer = Quantizer(shape=shape) def forward(self, x): return self.module(self.quantizer.quantize(x)) def add_actquant(module, name='', layers=[nn.Conv2d, nn.Linear]): if isinstance(module, ActQuantWrapper): return for attr in dir(module): tmp = getattr(module, attr) if type(tmp) in layers: setattr(module, attr, ActQuantWrapper(tmp)) if type(tmp) == nn.Sequential: replaced = [] for i, child in enumerate(tmp.children()): if type(child) in layers: replaced.append(ActQuantWrapper(child)) else: replaced.append(child) setattr(module, attr, nn.Sequential(*replaced)) if type(tmp) == torch.nn.ModuleList: replaced = [] for i, child in enumerate(tmp.children()): if type(child) in layers: replaced.append(ActQuantWrapper(child)) else: replaced.append(child) setattr(module, attr, nn.ModuleList(replaced)) for name1, child in module.named_children(): add_actquant(child, name + '.' + name1 if name != '' else name1, layers) import time class Quant4Linear(nn.Module): def __init__(self, linear, scales, zeros): super().__init__() self.register_buffer('zeros', zeros.clone() * scales) self.register_buffer('scales', scales) self.register_buffer('bias', linear.bias.data) intweight = torch.round((linear.weight.data + self.zeros) / self.scales).to(torch.int) intweight = intweight.t().contiguous() self.register_buffer('qweight', torch.zeros( (intweight.shape[0] // 8, intweight.shape[1]), dtype=torch.int, device=self.bias.device )) for i in range(intweight.shape[0]): self.qweight[i // 8] |= intweight[i] << (4 * (i % 8)) # self.linear = linear.to(torch.device('cuda:0')) def forward(self, x): if x.shape[-1] == x.numel(): outshape = list(x.shape) y = self.bias.clone() outshape[-1] = self.bias.numel() quant_cuda.vecquant4matmul(x, self.qweight, y, self.scales, self.zeros) # y = self.linear(x) return y.reshape(outshape) print(x.shape) raise ValueError('Only supports a single token currently.') def make_quant4(module, quantizers, name=''): if isinstance(module, Quant4Linear): return for attr in dir(module): tmp = getattr(module, attr) name1 = name + '.' + attr if name != '' else attr if name1 in quantizers: setattr( module, attr, Quant4Linear(tmp, quantizers[name1].scale, quantizers[name1].zero) ) for name1, child in module.named_children(): make_quant4(child, quantizers, name + '.' + name1 if name != '' else name1)