| import torch |
| from torch import nn |
| import quiptools_cuda |
|
|
| from lib.utils.matmul_had import matmul_hadU_cuda, matmul_hadUt_cuda |
|
|
|
|
| def get_grid(): |
| hintr = torch.arange(-8, 8) + 1 / 2 |
| return hintr.unsqueeze(-1) |
|
|
|
|
| _HI4B1C_CACHED = get_grid() |
| _HI4B1C_NORM_CACHED = torch.diag(_HI4B1C_CACHED @ _HI4B1C_CACHED.T) |
|
|
|
|
| class HI4B1C_codebook(nn.Module): |
|
|
| def __init__(self, inference=False): |
| super(HI4B1C_codebook, self).__init__() |
| self.opt_scale = 2.97 |
| self.codesz = 1 |
| self.idx_dtype = torch.int32 |
| self.packsz = 8 |
| self.pack_out = False |
| self.version = 0 |
|
|
| self.register_buffer('grid', _HI4B1C_CACHED) |
| if not inference: |
| self.register_buffer('grid_norm', _HI4B1C_NORM_CACHED) |
| ''' |
| self.cuda() |
| samples = torch.distributions.multivariate_normal.MultivariateNormal(torch.zeros(1), torch.eye(1)).rsample([200000]).cuda() |
| print(samples.shape) |
| def fn_s(s): |
| err = (self.quantize(samples*s, False)/s - samples).float().norm()**2 |
| err = err.cpu() / torch.numel(samples) |
| return err.cpu() |
| import scipy |
| print(scipy.optimize.minimize_scalar(fn_s, bounds=(0.1, 100))) |
| exit() |
| ''' |
|
|
| def round(self, X, grid, grid_norm): |
| assert X.shape[-1] == self.codesz |
| Xqidx = (2 * X @ grid.T - grid_norm).argmax(-1) |
| return grid[Xqidx], Xqidx |
|
|
| def quantize(self, X, return_idx=True): |
| vals, idx = self.round(X, self.grid, self.grid_norm) |
| if not return_idx: |
| return vals |
| return vals, idx.to(self.idx_dtype) |
|
|
| def maybe_pack_idxs(self, idxs): |
| return \ |
| (idxs[:, 0::self.packsz] << 4*7) + \ |
| (idxs[:, 2::self.packsz] << 4*6) + \ |
| (idxs[:, 4::self.packsz] << 4*5) + \ |
| (idxs[:, 6::self.packsz] << 4*4) + \ |
| (idxs[:, 1::self.packsz] << 4*3) + \ |
| (idxs[:, 3::self.packsz] << 4*2) + \ |
| (idxs[:, 5::self.packsz] << 4*1) + \ |
| idxs[:, 7::self.packsz] |
|
|
| def by_idxs(self, idxs, packed=False): |
| if packed: |
| idxs = idxs.repeat_interleave(self.packsz, dim=-1) |
| idxs[:, 0::self.packsz] = (idxs[:, 0::self.packsz] >> 28) & 15 |
| idxs[:, 2::self.packsz] = (idxs[:, 2::self.packsz] >> 24) & 15 |
| idxs[:, 4::self.packsz] = (idxs[:, 4::self.packsz] >> 20) & 15 |
| idxs[:, 6::self.packsz] = (idxs[:, 6::self.packsz] >> 16) & 15 |
| idxs[:, 1::self.packsz] = (idxs[:, 1::self.packsz] >> 12) & 15 |
| idxs[:, 3::self.packsz] = (idxs[:, 3::self.packsz] >> 8) & 15 |
| idxs[:, 5::self.packsz] = (idxs[:, 5::self.packsz] >> 4) & 15 |
| idxs[:, 7::self.packsz] = idxs[:, 7::self.packsz] & 15 |
|
|
| return self.grid[idxs.int()] |
|
|
|
|
| class QuantizedHI4B1CLinear(nn.Module): |
|
|
| def __init__(self, device): |
| super().__init__() |
| self.codebook = HI4B1C_codebook(inference=True).to(torch.float16).to(device) |
|
|
| def forward(self, |
| input, |
| Qidxs, |
| SU, |
| SV, |
| Wscale, |
| had_left, |
| had_right, |
| K_left, |
| K_right, |
| rank=-1, |
| A=None, |
| B=None, |
| rescale_WH=False, |
| scaleWH=None, |
| packed=False): |
| n, m = len(SU), len(SV) |
|
|
| x = input.view(-1, n).to(torch.float32) |
| if rescale_WH: |
| x /= scaleWH |
| x = x * SU |
| x = matmul_hadUt_cuda(x, had_left, K_left) |
|
|
| if rank > 0: |
| Bx = x @ B.t().to(torch.float32) |
| ABx = Bx @ A.t().to(torch.float32) |
|
|
| num_scale = 1024 |
| x = x / num_scale |
| x = x.to(torch.float16) |
|
|
| if packed: |
| W_decompressed = torch.zeros(m, n, dtype=torch.float16, device=x.device) |
| quiptools_cuda.decompress_hi4b1c_packed(Qidxs, self.codebook.grid, W_decompressed) |
| else: |
| W_decompressed = self.codebook.by_idxs(Qidxs, packed=False).reshape(-1, n) |
|
|
| z = x @ W_decompressed.t() |
|
|
| x = z.to(torch.float32) |
| x = x * (Wscale * num_scale) |
|
|
| if rank > 0: |
| x = x + ABx.to(torch.float32) |
|
|
| x = matmul_hadU_cuda(x, had_right, K_right) |
| x = x * SV |
|
|
| output = x.view(*input.shape[:-1], m) |
|
|
| return output |
|
|
