| 41 | |
| 42 | |
| 43 | class Quantizer: |
| 44 | |
| 45 | def __init__(self, config: Dict) -> None: |
| 46 | self.config = config |
| 47 | assert self.config['num_bits'] == 4 or self.config[ |
| 48 | 'num_bits'] == 8, 'Only INT4 and INT8 quantization is supported.' |
| 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 |
| 75 | min_value = tensor.amin(dim=self.config['group_dim'] + 1, keepdim=True) |
| 76 | max_value = tensor.amax(dim=self.config['group_dim'] + 1, keepdim=True) |
| 77 | |
| 78 | scale = q_range / (max_value - min_value) |
| 79 | |
| 80 | tensor = tensor.sub_(min_value).mul_(scale) |
| 81 | tensor = tensor_round(tensor_clamp(tensor, 0, q_range)).to(torch.uint8) |
| 82 | return tensor, scale, min_value |
| 83 | |
| 84 | def _compress_uint8_to_uint4(self, tensor: Tensor) -> Tensor: |
| 85 | assert tensor.shape[-1] % 2 == 0 |
| 86 | |
| 87 | new_data_shape = list(tensor.shape) |
| 88 | new_data_shape[-1] = new_data_shape[-1] // 2 |
| 89 | |
| 90 | data = torch.empty(new_data_shape, dtype=torch.uint8, device=tensor.device) |
| 91 | data = torch.bitwise_or(tensor[..., 0::2].bitwise_left_shift(4), tensor[..., 1::2]) |
| 92 | |
| 93 | return data |
| 94 | |
| 95 | |
| 96 | class DeQuantizer: |
no outgoing calls