| import torch |
| import torch.nn as nn |
| import quiptools_cuda |
| from lib.utils import dtype_from_str, get_hadK |
| from lib import codebook |
| import time |
|
|
|
|
| class QuantizedLinear(nn.Module): |
|
|
| def __init__(self, |
| in_features, |
| out_features, |
| codesz, |
| packsz, |
| pack_out, |
| idx_dtype, |
| codebook_version, |
| outlier_channel_split=False, |
| rank=-1, |
| rescale_WH=False, |
| bias=False): |
| super().__init__() |
|
|
| self.in_features = in_features |
| self.out_features = out_features |
| self.outlier_channel_split = outlier_channel_split |
| self.rank = rank |
| self.rescale_WH = rescale_WH |
|
|
| self.has_bias = bias |
| if self.has_bias: |
| self.register_buffer('bias', torch.ones(out_features)) |
| |
| if self.outlier_channel_split: |
| self.register_buffer('ocs_dupe_inds', torch.arange(in_features)) |
|
|
| if self.rank > 0: |
| self.register_buffer('A', torch.zeros(out_features, rank)) |
| self.register_buffer('B', torch.zeros(rank, in_features)) |
| else: |
| self.A = None |
| self.B = None |
|
|
| if self.rescale_WH: |
| self.register_buffer("scaleWH", torch.ones(in_features)) |
| else: |
| self.scaleWH = None |
|
|
| |
| if pack_out: |
| self.register_buffer( |
| "Qidxs", |
| torch.zeros(out_features // packsz, |
| in_features // codesz, |
| dtype=dtype_from_str(idx_dtype))) |
| else: |
| self.register_buffer( |
| "Qidxs", |
| torch.zeros(out_features, |
| in_features // (codesz * packsz), |
| dtype=dtype_from_str(idx_dtype))) |
|
|
| self.register_buffer("codebook_id", torch.tensor(0)) |
| self.register_buffer("SU", torch.ones(in_features)) |
| self.register_buffer("SV", torch.ones(out_features)) |
| self.register_buffer("Wscale", torch.ones(())) |
|
|
| self.built_codebook_class = False |
| self.built_graph = False |
| self.codebook_version = codebook_version |
|
|
| had_left, K_left = get_hadK(in_features) |
| had_right, K_right = get_hadK(out_features) |
| self.register_buffer('had_left', had_left, persistent=False) |
| self.register_buffer('had_right', had_right, persistent=False) |
| self.K_left = K_left |
| self.K_right = K_right |
| self.packed = (packsz != 1) |
|
|
| def forward(self, input): |
| if not self.built_codebook_class: |
| self.codebook_class = codebook.get_quantized_class(self.codebook_id.item())( |
| self.Qidxs.device) |
| if self.codebook_class.codebook.version != self.codebook_version: |
| raise Exception( |
| f"Saved weights version ({self.codebook_version}) does not match the "\ |
| f"codebook version ({self.codebook_class.codebook.version}). "\ |
| "Please download the latest weights from https://huggingface.co/relaxml") |
| self.built_codebook_class = True |
|
|
| if self.outlier_channel_split: |
| input = input[..., self.ocs_dupe_inds] |
|
|
| result = self.codebook_class(input, |
| self.Qidxs, |
| self.SU, |
| self.SV, |
| self.Wscale, |
| self.had_left, |
| self.had_right, |
| self.K_left, |
| self.K_right, |
| rank=self.rank, |
| A=self.A, |
| B=self.B, |
| rescale_WH=self.rescale_WH, |
| scaleWH=self.scaleWH, |
| packed=self.packed) |
| if self.has_bias: |
| return result + self.bias |
| return result |
| |
|
|
