| 21 | |
| 22 | |
| 23 | class UniformAffineQuantizer(nn.Module): |
| 24 | def __init__( |
| 25 | self, |
| 26 | n_bits: int = 8, |
| 27 | group_size=None, |
| 28 | weight=None, |
| 29 | ): |
| 30 | super().__init__() |
| 31 | assert 2 <= n_bits <= 16, "bitwidth not supported" |
| 32 | self.n_bits = n_bits |
| 33 | self.qmin = 0 |
| 34 | self.qmax = 2 ** (n_bits) - 1 |
| 35 | self.group_size = group_size if group_size != -1 else weight.shape[-1] |
| 36 | assert weight.shape[-1] % group_size == 0 |
| 37 | self.enable = True |
| 38 | |
| 39 | # init scale and zero point through Max-Min quantization |
| 40 | with torch.no_grad(): |
| 41 | if weight is not None: |
| 42 | x = weight.reshape(-1,self.group_size) |
| 43 | xmin = x.amin([-1], keepdim=True) |
| 44 | xmax = x.amax([-1], keepdim=True) |
| 45 | range = xmax - xmin |
| 46 | scale = range / (2**self.n_bits-1) |
| 47 | scale = scale.clamp(min=1e-4, max=1e4) |
| 48 | zero_point = -(xmin/scale).clamp(min=-1e4, max=1e4) |
| 49 | self.scale = nn.Parameter(scale) |
| 50 | self.zero_point = nn.Parameter(zero_point.round()) |
| 51 | |
| 52 | |
| 53 | def change_n_bits(self, n_bits): |
| 54 | self.n_bits = n_bits |
| 55 | self.qmin = 0 |
| 56 | self.qmax = int(2 ** (n_bits) - 1) |
| 57 | |
| 58 | def fake_quant(self, x): |
| 59 | scale = clamp_ste(self.scale,1e-4, 1e4) |
| 60 | round_zero_point = clamp_ste(round_ste(self.zero_point), self.qmin, self.qmax) |
| 61 | |
| 62 | dim1, dim2 = x.shape |
| 63 | x = x.reshape(-1, self.group_size) |
| 64 | x_int = round_ste(x / scale) |
| 65 | if round_zero_point is not None: |
| 66 | x_int = x_int.add(round_zero_point) |
| 67 | x_int = x_int.clamp(self.qmin, self.qmax) |
| 68 | x_dequant = x_int |
| 69 | if round_zero_point is not None: |
| 70 | x_dequant = x_dequant.sub(round_zero_point) |
| 71 | x_dequant = x_dequant.mul(scale) |
| 72 | if self.group_size: |
| 73 | x_dequant = x_dequant.reshape(dim1, dim2) |
| 74 | return x_dequant |
| 75 | |
| 76 | |
| 77 | def forward(self, x: torch.Tensor): |
| 78 | if self.n_bits >= 16 or not self.enable: |
| 79 | return x |
| 80 | |