(cls, linear, w_bit, group_size, init_only=False, scales=None, zeros=None)
| 69 | |
| 70 | @classmethod |
| 71 | def from_linear(cls, linear, w_bit, group_size, init_only=False, scales=None, zeros=None): |
| 72 | awq_linear = cls(w_bit, group_size, linear.in_features, linear.out_features, linear.bias is not None, linear.weight.device) |
| 73 | if init_only: # just prepare for loading sd |
| 74 | return awq_linear |
| 75 | |
| 76 | # need scales and zeros info for real quantization |
| 77 | assert scales is not None and zeros is not None |
| 78 | scale_zeros = zeros * scales |
| 79 | |
| 80 | pack_num = 32 // awq_linear.w_bit |
| 81 | |
| 82 | qscales = torch.zeros( |
| 83 | (scales.shape[0], calculate_zeros_width(linear.in_features, group_size, pack_num) * pack_num), |
| 84 | dtype=torch.float16, |
| 85 | device=scales.device |
| 86 | ) |
| 87 | qscales[:, :scales.shape[1]] = scales |
| 88 | # awq_linear.scales = scales.clone().half() |
| 89 | awq_linear.scales = qscales |
| 90 | |
| 91 | if linear.bias is not None: |
| 92 | awq_linear.bias = linear.bias.clone().half() |
| 93 | |
| 94 | intweight = [] |
| 95 | for idx in range(awq_linear.in_features): |
| 96 | intweight.append(torch.round((linear.weight.data[:, idx] + scale_zeros[:, idx // group_size]) / awq_linear.scales[:, idx // group_size]).to(torch.int)[:, None]) |
| 97 | intweight = torch.cat(intweight, dim=1) |
| 98 | # intweight = intweight.t().contiguous() |
| 99 | intweight = intweight.to(dtype=torch.int32) |
| 100 | |
| 101 | qweight = torch.zeros((intweight.shape[0], intweight.shape[1] // 32 * awq_linear.w_bit), dtype=torch.int32, device=intweight.device) |
| 102 | |
| 103 | for col in range(intweight.shape[1] // pack_num): |
| 104 | if awq_linear.w_bit == 4: |
| 105 | # order_map = [0, 2, 4, 6, 1, 3, 5, 7] |
| 106 | order_map = [0, 1, 2, 3, 4, 5, 6, 7] |
| 107 | elif awq_linear.w_bit == 2: |
| 108 | import numpy as np |
| 109 | order_map = np.arange(16) |
| 110 | else: |
| 111 | raise NotImplementedError("Only 4-bit are supported for now.") |
| 112 | for i in range(pack_num): |
| 113 | qweight_col = intweight[:, col * pack_num + order_map[i]] |
| 114 | qweight[:, col] |= qweight_col << (i * awq_linear.w_bit) |
| 115 | |
| 116 | awq_linear.qweight = qweight |
| 117 | |
| 118 | zeros = zeros.to(dtype=torch.int32) |
| 119 | qzeros = torch.zeros( |
| 120 | (zeros.shape[0], calculate_zeros_width(linear.in_features, group_size, pack_num)), |
| 121 | dtype=torch.int32, |
| 122 | device=zeros.device, |
| 123 | ) |
| 124 | |
| 125 | for col in range((zeros.shape[1] + pack_num - 1) // pack_num): |
| 126 | if awq_linear.w_bit == 4: |
| 127 | # order_map = [0, 2, 4, 6, 1, 3, 5, 7] |
| 128 | order_map = [0, 1, 2, 3, 4, 5, 6, 7] |
no test coverage detected