(self, tensor: Tensor)
| 49 | assert self.config['symmetric'] == False, 'Only asymmetric quantization is supported at this moment.' |
| 50 | |
| 51 | def quantize(self, tensor: Tensor) -> Tuple[Tensor, Tensor, Tensor]: |
| 52 | assert tensor.shape[self.config['group_dim']] % self.config['group_size'] == 0 \ |
| 53 | , f'Tensor shape: {tensor.shape} quantization config {self.config}' |
| 54 | |
| 55 | tensor = torch.clone(tensor) |
| 56 | |
| 57 | shape = tensor.shape |
| 58 | num_groups = shape[self.config['group_dim']] // self.config['group_size'] |
| 59 | new_shape = (shape[:self.config['group_dim']] + (num_groups, self.config['group_size']) + |
| 60 | shape[self.config['group_dim'] + 1:]) |
| 61 | tensor = tensor.view(new_shape) |
| 62 | |
| 63 | quantized_tensor, scale, min_value = self._quantize_int8(tensor) |
| 64 | quantized_tensor = quantized_tensor.view(shape) |
| 65 | |
| 66 | if self.config['num_bits'] == 4: |
| 67 | return self._compress_uint8_to_uint4(quantized_tensor), scale, min_value |
| 68 | if self.config['num_bits'] == 8: |
| 69 | return quantized_tensor, scale, min_value |
| 70 | |
| 71 | assert False, 'Unsupported quantization bits {}'.format(self.config['num_bits']) |
| 72 | |
| 73 | def _quantize_int8(self, tensor: Tensor) -> Tuple[Tensor, Tensor, Tensor]: |
| 74 | q_range = 2**self.config['num_bits'] - 1 |
no test coverage detected