| 10 | init(autoreset=True) |
| 11 | |
| 12 | class QuantLinear(nn.Module): |
| 13 | def __init__(self, in_features, out_features, groupsize=-1, double_groupsize=-1, bits=4, v1 = True, asym=True): |
| 14 | super().__init__() |
| 15 | |
| 16 | bits=bits |
| 17 | self.in_features = in_features |
| 18 | self.out_features = out_features |
| 19 | self.bits=bits |
| 20 | self.maxq = 2 ** self.bits - 1 |
| 21 | groupsize = groupsize if groupsize != -1 else in_features |
| 22 | self.groupsize = groupsize |
| 23 | |
| 24 | if double_groupsize==-1: |
| 25 | double_groupsize=out_features |
| 26 | |
| 27 | self.asym = asym |
| 28 | |
| 29 | self.disable_bias = True |
| 30 | self.initialize(in_features, out_features, groupsize, double_groupsize, bits, v1, asym) |
| 31 | |
| 32 | def initialize(self, in_features, out_features, groupsize, double_quantize_groupsize, bits, v1, asym): |
| 33 | |
| 34 | if asym: |
| 35 | self.register_buffer('qzeros', torch.empty((math.ceil(in_features/groupsize), math.ceil(out_features / 256 * (bits * 8))), dtype=torch.int32)) |
| 36 | if bits == 4: |
| 37 | self.register_buffer('qscales', torch.empty((math.ceil(in_features/groupsize), math.ceil(out_features/double_quantize_groupsize), double_quantize_groupsize), dtype=torch.uint8)) |
| 38 | else: |
| 39 | self.register_buffer('qscales', torch.empty((math.ceil(in_features/groupsize), out_features), dtype=torch.uint8)) |
| 40 | |
| 41 | else: |
| 42 | self.register_buffer('qstatistic', torch.empty((math.ceil(in_features/groupsize), math.ceil(out_features/double_quantize_groupsize), double_quantize_groupsize), dtype=torch.uint8)) |
| 43 | self.register_buffer('qzeros_zeros', torch.empty((math.ceil(in_features/groupsize), math.ceil(out_features/double_quantize_groupsize), 1), dtype=torch.half)) |
| 44 | self.register_buffer('qzeros_scales', torch.empty((math.ceil(in_features/groupsize), math.ceil(out_features/double_quantize_groupsize), 1), dtype=torch.half)) |
| 45 | |
| 46 | if not v1: |
| 47 | self.register_buffer('qscales_zeros', torch.empty((math.ceil(in_features/groupsize), math.ceil(out_features/double_quantize_groupsize), 1), dtype=torch.half)) |
| 48 | self.register_buffer('qscales_scales', torch.empty((math.ceil(in_features/groupsize), math.ceil(out_features/double_quantize_groupsize), 1), dtype=torch.half)) |
| 49 | else: |
| 50 | self.register_buffer('qscales_zeros', torch.empty((1, out_features, 1), dtype=torch.half)) |
| 51 | self.register_buffer('qscales_scales', torch.empty((1, out_features, 1), dtype=torch.half)) |
| 52 | |
| 53 | self.register_buffer('g_idx', torch.tensor([i // groupsize for i in range(in_features)], dtype=torch.int32)) |
| 54 | self.register_buffer('qweight', torch.empty(math.ceil(in_features / 256 * (bits * 8)), out_features, dtype=torch.int32)) |
| 55 | self.register_buffer("wf", torch.tensor(list(range(0,32,bits)), dtype=torch.int32).unsqueeze(0)) |
| 56 | self.register_buffer('bias', torch.empty(out_features)) |
| 57 | |
| 58 | def forward(self, x): |
| 59 | if self.bits in [2, 4, 8, 16]: |
| 60 | weight_unpack = torch.bitwise_right_shift(torch.unsqueeze(self.qweight, 1).expand(-1, 32 // self.bits, -1), self.wf.unsqueeze(-1)).to(torch.int16 if self.bits == 8 else torch.int8).view(-1, self.qweight.size(-1)) |
| 61 | torch.bitwise_and(weight_unpack,(2 ** self.bits) - 1, out=weight_unpack) |
| 62 | else: |
| 63 | raise ValueError |
| 64 | |
| 65 | if self.asym: |
| 66 | zeros_unpack = torch.bitwise_right_shift(torch.unsqueeze(self.qzeros, 2).expand(-1, -1, 32 // self.bits), self.wf.unsqueeze(0)).to(torch.int16 if self.bits == 8 else torch.int8) |
| 67 | torch.bitwise_and(zeros_unpack, (2 ** self.bits) - 1, out=zeros_unpack) |
| 68 | |
| 69 | zeros_unpack = zeros_unpack + 1 |