| 97 | del self.g_idx |
| 98 | |
| 99 | def pack(self, linear, scales, zeros, g_idx=None): |
| 100 | W = linear.weight.data.clone() |
| 101 | if isinstance(linear, nn.Conv2d): |
| 102 | W = W.flatten(1) |
| 103 | if isinstance(linear, transformers.pytorch_utils.Conv1D): |
| 104 | W = W.t() |
| 105 | |
| 106 | g_idx = torch.tensor([i // self.group_size for i in range(self.infeatures)], dtype=torch.int32) |
| 107 | |
| 108 | scale_zeros = zeros * scales |
| 109 | self.scales = nn.Parameter(scales.half()) |
| 110 | if linear.bias is not None: |
| 111 | self.bias = linear.bias.clone().half() |
| 112 | |
| 113 | intweight = [] |
| 114 | for idx in range(self.infeatures): |
| 115 | intweight.append( |
| 116 | torch.round( |
| 117 | ( |
| 118 | W[:, idx] + scale_zeros[g_idx[idx]]) / self.scales[g_idx[idx]] |
| 119 | ).to(torch.int)[:, None] |
| 120 | ) |
| 121 | intweight = torch.cat(intweight, dim=1) |
| 122 | intweight = intweight.t().contiguous() |
| 123 | intweight = intweight.numpy().astype(np.uint32) |
| 124 | |
| 125 | i = 0 |
| 126 | row = 0 |
| 127 | qweight = np.zeros((math.ceil(intweight.shape[0]/(32//self.bits)), intweight.shape[1]), dtype=np.uint32) |
| 128 | while row < qweight.shape[0]: |
| 129 | if self.bits in [2, 3, 4, 8]: |
| 130 | for j in range(i, min(i + (32 // self.bits), intweight.shape[0])): |
| 131 | qweight[row] |= intweight[j] << (self.bits * (j - i)) |
| 132 | i += 32 // self.bits |
| 133 | row += 1 |
| 134 | else: |
| 135 | raise NotImplementedError("Only 2,3,4,8 bits are supported.") |
| 136 | |
| 137 | qweight = qweight.astype(np.int32) |
| 138 | self.qweight = torch.from_numpy(qweight) |
| 139 | |
| 140 | zeros = zeros.numpy().astype(np.uint32) |
| 141 | self.zeros_dim0, self.zeros_dim1 = zeros.shape |
| 142 | qzeros = np.zeros((zeros.shape[0], math.ceil(zeros.shape[1] / (32 // self.bits))), dtype=np.uint32) |
| 143 | i = 0 |
| 144 | col = 0 |
| 145 | while col < qzeros.shape[1]: |
| 146 | if self.bits in [2, 3, 4, 8]: |
| 147 | for j in range(i, min(i + (32 // self.bits), zeros.shape[1])): |
| 148 | qzeros[:, col] |= zeros[:, j] << (self.bits * (j - i)) |
| 149 | i += 32 // self.bits |
| 150 | col += 1 |
| 151 | else: |
| 152 | raise NotImplementedError("Only 2,3,4,8 bits are supported.") |
| 153 | |
| 154 | qzeros = qzeros.astype(np.int32) |
| 155 | self.qzeros = torch.from_numpy(qzeros) |
| 156 | |