matq.py

import math
import time
import torch
import torch.nn as nn
import transformers
from loguru import logger
from quant import quantize


DEBUG = False 

torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False

class TensorQ:
    def __init__(self, layer, rank=32):
        self.layer = layer
        self.dev = self.layer.weight.device
        W = layer.weight.data.clone()
        if isinstance(self.layer, nn.Conv2d):
            W = W.flatten(1)
        if isinstance(self.layer, transformers.Conv1D):
            W = W.t()
        self.rank = rank
        self.decompose()
        self.rows = W.shape[0]
        self.columns = W.shape[1]
        self.L_columns = self.L.shape[1]
        self.H = torch.zeros((self.columns, self.columns), device=self.dev)
        self.H_R = torch.zeros((self.columns, self.columns), device=self.dev)
        self.H_L = torch.zeros((self.L_columns, self.L_columns), device=self.dev)
        self.nsamples = 0

    def add_batch_lr(self, inp, out):
        if DEBUG:
            self.inp1 = inp
            self.out1 = out
        if len(inp.shape) == 2:
            inp = inp.unsqueeze(0)
        tmp = inp.shape[0]
        if isinstance(self.layer, nn.Linear) or isinstance(self.layer, transformers.Conv1D):
            if len(inp.shape) == 3:
                inp = inp.reshape((-1, inp.shape[-1]))
            inp = inp.t()
        if isinstance(self.layer, nn.Conv2d):
            unfold = nn.Unfold(
                self.layer.kernel_size,
                dilation=self.layer.dilation,
                padding=self.layer.padding,
                stride=self.layer.stride
            )
            inp = unfold(inp)
            inp = inp.permute([1, 0, 2])
            inp = inp.flatten(1)
        self.H_R *= self.nsamples / (self.nsamples + tmp)
        self.nsamples += tmp
        inp = math.sqrt(2 / self.nsamples) * inp.float()

        self.H_R += inp.matmul(inp.t())
        # logger.info(f"self.H_R: {self.H_R.shape}")
        # for L, consider the input to be R@X
        inp = self.R @ inp
        self.H_L *= self.nsamples / (self.nsamples + tmp)
        self.H_L += inp.matmul(inp.t())
        # logger.info(f"self.H_L: {self.H_L.shape}")

    def free(self):
        if DEBUG:
            self.inp1 = None
            self.out1 = None
        self.H = None
        self.H_L = None
        self.H_R = None
        self.Losses = None
        self.Trace = None
        torch.cuda.empty_cache()

    def decompose(self):
        W = self.layer.weight.data.clone()
        if isinstance(self.layer, nn.Conv2d):
            W = W.flatten(1)
        if isinstance(self.layer, transformers.Conv1D):
            W = W.t()
        W = W.float()
        logger.info("starting decomposition")
        tick = time.time()
        U, S, Vh = torch.pca_lowrank(W, q=self.rank, center=True, niter=5)
        # let's say L = U
        # and R = diag(S)*V.T
        self.L = U
        self.R = torch.diag_embed(S) @ Vh.T
        logger.info(f"decomposition done. elapsed time: {time.time() - tick}, L: {self.L.shape}, R: {self.R.shape}")

    def lr_quant(self, blocksize=128, percdamp=.01, groupsize=-1, actorder=False):
        self.lr_quant_R(blocksize, percdamp, groupsize, actorder)
        self.lr_quant_L(blocksize, percdamp, groupsize, actorder)
        # restored weight is L@R
        # but on disk we only save L, R
        self.layer.weight.data = (self.L @ self.R).reshape(self.layer.weight.shape).to(self.layer.weight.data.dtype)
        
    def lr_quant_R(self, blocksize=128, percdamp=.01, groupsize=-1, actorder=False):
        R = self.R.data.clone()
        if isinstance(self.layer, nn.Conv2d):
            R = R.flatten(1)
        if isinstance(self.layer, transformers.Conv1D):
            R = R.t()
        R = R.float()

        tick = time.time()
        # now quantize R first
        # logger.info("quantizing R")
        if not self.quantizer.ready():
            self.quantizer.find_params(R, weight=True)
        H_R = self.H_R
        del self.H_R
        dead = torch.diag(H_R) == 0
        H_R[dead, dead] = 1
        R[:, dead] = 0

        if actorder:
            perm = torch.argsort(torch.diag(H_R), descending=True)
            R = R[:, perm]
            H_R = H_R[perm][:, perm]
        
        Losses_R = torch.zeros_like(R)
        Q_R = torch.zeros_like(R)

        damp = percdamp * torch.mean(torch.diag(H_R))
        diag = torch.arange(self.columns, device=self.dev)
        H_R[diag, diag] += damp

        H_R = torch.linalg.cholesky(H_R)
        H_R = torch.cholesky_inverse(H_R)
        H_R = torch.linalg.cholesky(H_R, upper=True)

        Hinv_R = H_R

        for i1 in range(0, self.columns, blocksize):
            i2 = min(i1 + blocksize, self.columns)
            count = i2 - i1

            W1 = R[:, i1:i2].clone()
            Q1 = torch.zeros_like(W1)
            Err1 = torch.zeros_like(W1)
            Losses1 = torch.zeros_like(W1)
            Hinv1 = Hinv_R[i1:i2, i1:i2]

            for i in range(count):
                w = W1[:, i]
                d = Hinv1[i, i]
                if groupsize != -1:
                    if (i1 + i) % groupsize == 0:
                        self.quantizer.find_params(R[:, (i1 + i):(i1 + i + groupsize)], weight=True)
                
                q = quantize(
                    w.unsqueeze(1), self.quantizer.scale, self.quantizer.zero, self.quantizer.maxq
                ).flatten()
                q
                Q1[:, i] = q
                Losses1[:, i] = (w - q) ** 2 / d ** 2

                err1 = (w - q) / d

                W1[:, i:] -= err1.unsqueeze(1).matmul(Hinv1[i, i:].unsqueeze(0))

                Err1[:, i] = err1
            
            Q_R[:, i1:i2] = Q1
            Losses_R[:, i1:i2] = Losses1 / 2

            R[:, i2:] -= Err1.matmul(Hinv_R[i1:i2, i2:])
            if DEBUG:
                self.layer.weight.data[:, :i2] = Q_R[:, :i2]
                self.layer.weight.data[:, i2:] = R[:, i2:]
                print(torch.sum((self.layer(self.inp1) - self.out1) ** 2))
                print(torch.sum(Losses_R))
        torch.cuda.synchronize()
        total_time = time.time() - tick
        error = torch.sum(Losses_R).item()
        if actorder:
            invperm = torch.argsort(perm)
            Q_R = Q_R[:, invperm]
        
        if isinstance(self.layer, transformers.Conv1D):
            Q_R = Q_R.t()
        self.R = Q_R.reshape(self.R.shape).to(self.R.dtype)

    def lr_quant_L(self, blocksize=128, percdamp=.01, groupsize=-1, actorder=False):
        L = self.L.data.clone()
        if isinstance(self.layer, nn.Conv2d):
            L = L.flatten(1)
        if isinstance(self.layer, transformers.Conv1D):
            L = L.t()
        L = L.float()
        tick = time.time()
        # now quantize R first
        # logger.info("quantizing L")
        if not self.l_quantizer.ready():
            self.l_quantizer.find_params(L, weight=True)
        H_L = self.H_L
        del self.H_L
        dead = torch.diag(H_L) == 0
        H_L[dead, dead] = 1
        L[:, dead] = 0

        if actorder:
            perm = torch.argsort(torch.diag(H_L), descending=True)
            L = L[:, perm]
            H_L = H_L[perm][:, perm]
        
        Losses_L = torch.zeros_like(L)
        Q_L = torch.zeros_like(L)

        damp = percdamp * torch.mean(torch.diag(H_L))
        diag = torch.arange(self.L_columns, device=self.dev)
        H_L[diag, diag] += damp

        H_L = torch.linalg.cholesky(H_L)
        H_L = torch.cholesky_inverse(H_L)
        H_L = torch.linalg.cholesky(H_L, upper=True)

        Hinv_L = H_L

        for i1 in range(0, self.L_columns, blocksize):
            i2 = min(i1 + blocksize, self.L_columns)
            count = i2 - i1

            W1 = L[:, i1:i2].clone()
            Q1 = torch.zeros_like(W1)
            Err1 = torch.zeros_like(W1)
            Losses1 = torch.zeros_like(W1)
            Hinv1 = Hinv_L[i1:i2, i1:i2]

            for i in range(count):
                w = W1[:, i]
                d = Hinv1[i, i]
                if groupsize != -1:
                    if (i1 + i) % groupsize == 0:
                        self.l_quantizer.find_params(L[:, (i1 + i):(i1 + i + groupsize)], weight=True)
                
                q = quantize(
                    w.unsqueeze(1), self.l_quantizer.scale, self.l_quantizer.zero, self.l_quantizer.maxq
                ).flatten()

                Q1[:, i] = q
                Losses1[:, i] = (w - q) ** 2 / d ** 2

                err1 = (w - q) / d

                W1[:, i:] -= err1.unsqueeze(1).matmul(Hinv1[i, i:].unsqueeze(0))

                Err1[:, i] = err1
            
            Q_L[:, i1:i2] = Q1
            Losses_L[:, i1:i2] = Losses1 / 2

            L[:, i2:] -= Err1.matmul(Hinv_L[i1:i2, i2:])
            if DEBUG:
                self.layer.weight.data[:, :i2] = Q_L[:, :i2]
                self.layer.weight.data[:, i2:] = L[:, i2:]
                print(torch.sum((self.layer(self.inp1) - self.out1) ** 2))
                print(torch.sum(Losses_L))
        torch.cuda.synchronize()
        total_time = time.time() - tick
        error = torch.sum(Losses_L).item()
        if actorder:
            invperm = torch.argsort(perm)
            Q_L = Q_L[:, invperm]
        
        if isinstance(self.layer, transformers.Conv1D):
            Q_L = Q_L.t()
        self.L = Q_L.reshape(self.L.shape).to(self.L.dtype)