| 293 | return grad_input, None, None, None, None, None, None |
| 294 | |
| 295 | class QuantLinear(nn.Module): |
| 296 | def __init__(self, bits, groupsize, infeatures, outfeatures, bias): |
| 297 | super().__init__() |
| 298 | if bits not in [2, 4, 8]: |
| 299 | raise NotImplementedError("Only 2,4,8 bits are supported.") |
| 300 | self.infeatures = infeatures |
| 301 | self.outfeatures = outfeatures |
| 302 | self.bits = bits |
| 303 | self.maxq = 2 ** self.bits - 1 |
| 304 | self.groupsize = groupsize if groupsize != -1 else infeatures |
| 305 | |
| 306 | self.register_buffer('qweight', torch.zeros((infeatures // 32 * self.bits, outfeatures), dtype=torch.int32)) |
| 307 | self.register_buffer('qzeros', torch.zeros((math.ceil(infeatures / self.groupsize), outfeatures // 32 * self.bits), dtype=torch.int32)) |
| 308 | self.register_buffer('scales', torch.zeros((math.ceil(infeatures / self.groupsize), outfeatures), dtype=torch.float16)) |
| 309 | self.register_buffer('g_idx', torch.tensor([i // self.groupsize for i in range(infeatures)], dtype=torch.int32)) |
| 310 | if bias: |
| 311 | self.register_buffer('bias', torch.zeros((outfeatures), dtype=torch.float16)) |
| 312 | else: |
| 313 | self.bias = None |
| 314 | |
| 315 | def pack(self, linear, scales, zeros, g_idx=None): |
| 316 | self.g_idx = g_idx.clone() if g_idx is not None else self.g_idx |
| 317 | |
| 318 | scales = scales.t().contiguous() |
| 319 | zeros = zeros.t().contiguous() |
| 320 | scale_zeros = zeros * scales |
| 321 | self.scales = scales.clone().half() |
| 322 | if linear.bias is not None: |
| 323 | self.bias = linear.bias.clone().half() |
| 324 | |
| 325 | intweight = [] |
| 326 | for idx in range(self.infeatures): |
| 327 | intweight.append(torch.round( |
| 328 | (linear.weight.data[:, idx] + scale_zeros[self.g_idx[idx]]) / self.scales[self.g_idx[idx]]).to( |
| 329 | torch.int)[:, None]) |
| 330 | intweight = torch.cat(intweight, dim=1) |
| 331 | intweight = intweight.t().contiguous() |
| 332 | intweight = intweight.numpy().astype(np.uint32) |
| 333 | qweight = np.zeros((intweight.shape[0] // 32 * self.bits, intweight.shape[1]), dtype=np.uint32) |
| 334 | i = 0 |
| 335 | row = 0 |
| 336 | while row < qweight.shape[0]: |
| 337 | if self.bits in [2, 4, 8]: |
| 338 | for j in range(i, i + (32 // self.bits)): |
| 339 | qweight[row] |= intweight[j] << (self.bits * (j - i)) |
| 340 | i += 32 // self.bits |
| 341 | row += 1 |
| 342 | else: |
| 343 | raise NotImplementedError("Only 2,4,8 bits are supported.") |
| 344 | |
| 345 | qweight = qweight.astype(np.int32) |
| 346 | self.qweight = torch.from_numpy(qweight) |
| 347 | |
| 348 | zeros -= 1 |
| 349 | zeros = zeros.numpy().astype(np.uint32) |
| 350 | qzeros = np.zeros((zeros.shape[0], zeros.shape[1] // 32 * self.bits), dtype=np.uint32) |
| 351 | i = 0 |
| 352 | col = 0 |